[Bugfix][VM] Properly convert tensor inputs in save_function (#257)

It was observed that closures saved using `save_function` would crash when used over RPC with the `time_evaluator`, whereas using `set_input` and `invoke_stateful` worked as normal. While I am not entirely sure why these failures happened over RPC only in `time_evaluator` (but not in other RPC trials), it became clear that `set_input` performs a conversion of input tensor values in `SetInputTensorWithIndex`, while `save_function` was not doing this. Adding this conversion fixed the observed bug.

Author: slyubomirsky

src/runtime/relax_vm/vm.cc

@@ -127,7 +127,7 @@ PackedFunc VirtualMachine::GetFunction(const std::string& name,
       if (args.size() > 3) {
         inputs = std::vector<RegType>(args.size() - 3);
         for (int i = 3; i < args.size(); i++) {
-          inputs[i - 3] = args[i];
+          SetInputTensorWithIndex(inputs, args[i], i - 3, devices[0]);
         }
       }
       if (include_return) {

tests/python/relax/test_vm.py

@@ -1208,6 +1208,24 @@ def test_save_function_kwargs_rpc():
     run_on_rpc(TestVMSetInput, save_function_kwargs_trial)
 
 
+def save_function_time_evaluator_trial(
+    vm: relax.VirtualMachine, device: tvm.runtime.Device
+) -> None:
+    # just checking that the saved function can be called in the time evaluator
+    a = tvm.nd.array(np.random.rand(32, 32).astype("float32"), device)
+    b = tvm.nd.array(np.random.rand(32, 32).astype("float32"), device)
+    vm.save_function("main", "saved_main", a, b)
+    vm.time_evaluator("saved_main", device)()
+
+
+def test_save_function_time_evaluator():
+    save_function_time_evaluator_trial(*make_vm(TestVMSetInput))
+
+
+def test_save_function_time_evaluator():
+    run_on_rpc(TestVMSetInput, save_function_time_evaluator_trial)
+
+
 # if you set an input, you should not be able to call statelessly
 @pytest.mark.xfail()
 def test_set_input_stateless_failure():