diff --git a/src/relay/backend/te_compiler.cc b/src/relay/backend/te_compiler.cc index af243c02674e..c2352311a54d 100644 --- a/src/relay/backend/te_compiler.cc +++ b/src/relay/backend/te_compiler.cc @@ -726,19 +726,23 @@ LoweredModule LowerTE(const IRModule& module, TargetMap targets, DeviceMap devic auto updated_module = pass(module); - const auto* te_compiler_update_weights = - runtime::Registry::Get("auto_scheduler.relay_integration.te_compiler_update_weights"); + // A temporary solution until we can rewrite the auto-scheduler task extraction code to work + // in a more reasonable way. + if (backend::IsAutoSchedulerEnabled()) { + const auto* te_compiler_update_weights = + runtime::Registry::Get("auto_scheduler.relay_integration.te_compiler_update_weights"); - ICHECK(te_compiler_update_weights != nullptr) - << "auto_scheduler.relay_integration.te_compiler_update_weights"; + ICHECK(te_compiler_update_weights != nullptr) + << "auto_scheduler.relay_integration.te_compiler_update_weights"; - Map weight_map; + Map weight_map; - for (auto pair : compiler->GetOpWeights()) { - weight_map.Set(pair.first, pair.second); - } + for (auto pair : compiler->GetOpWeights()) { + weight_map.Set(pair.first, pair.second); + } - (*te_compiler_update_weights)(weight_map); + (*te_compiler_update_weights)(weight_map); + } LoweredModule lowered_module; lowered_module.main_module = updated_module;