zhiics commented on a change in pull request #6337:
URL: https://github.com/apache/incubator-tvm/pull/6337#discussion_r479556767



##########
File path: src/runtime/vm/vm.cc
##########
@@ -146,12 +155,17 @@ PackedFunc VirtualMachine::GetFunction(const std::string& 
name,
       auto func_index = gvit->second;
       const auto& vm_func = exec_->functions[func_index];
       const auto& param_names = vm_func.params;
-      // TODO(icemelon9): For heterogeneous execution, get input device 
information
-      TVMContext ctx = ctxs_[0];
       CHECK_EQ(args.size() - 1, param_names.size())
           << "The number of provided parameters doesn't match the number of 
arguments";
+      CHECK_EQ(param_names.size(), vm_func.params_device_type.size())
+          << "The number of provided parameters doesn't match the number of 
assigned devices";
       std::vector<ObjectRef> func_args(param_names.size());
       for (int i = 1; i < args.size(); ++i) {
+        TVMContext ctx;
+        int device_type = vm_func.params_device_type[i - 1];
+        ctx.device_type = DLDeviceType(device_type);

Review comment:
       good point




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
[email protected]


Reply via email to