[relay] Small refactor for context (#4091)

diff --git a/src/relay/backend/build_module.cc b/src/relay/backend/build_module.cc
index 20e760f..4cf13a3 100644 (file)
--- a/src/relay/backend/build_module.cc
+++ b/src/relay/backend/build_module.cc
@@ -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);
     }

diff --git a/src/relay/backend/vm/compiler.cc b/src/relay/backend/vm/compiler.cc
index 49079fb..013d8b0 100644 (file)
--- a/src/relay/backend/vm/compiler.cc
+++ b/src/relay/backend/vm/compiler.cc
@@ -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);
 }