Commit c12275ee by Zhi Committed by Tianqi Chen

[relay] Small refactor for context (#4091)

parent 3a32729c
......@@ -328,10 +328,9 @@ class RelayBuildModule : public runtime::ModuleNode {
// Create a sequential pass and perform optimizations.
transform::Pass seq = transform::Sequential(pass_seqs);
if (targets.size() == 1) {
for (const auto& kv : targets) {
With<Target> tctx(kv.second);
relay_module = seq(relay_module);
}
const auto& it = targets.begin();
With<Target> tctx((*it).second);
relay_module = seq(relay_module);
} else {
relay_module = seq(relay_module);
}
......
......@@ -897,10 +897,9 @@ Module VMCompiler::OptimizeModule(const Module& mod, const TargetsMap& targets)
// TODO(wweic): Support heterogenous execution
tvm::With<relay::transform::PassContext> ctx(pass_ctx);
if (targets.size() == 1) {
for (const auto& kv : targets) {
With<Target> tctx(kv.second);
return seq(mod);
}
const auto& it = targets.begin();
With<Target> tctx((*it).second);
return seq(mod);
}
return seq(mod);
}
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment