fix

Author: fangfangssj

graph_net/torch/backend/tvm_backend.py

@@ -33,9 +33,9 @@ def forward(self, **kwargs):
                     param = kwargs[name]
                     self.tvm_input.append(tvm.nd.array(param.cpu(), self.dev))
 
-        output = self.compiled_vm["subgraph_0"](*self.tvm_input).numpy()
+        output = self.compiled_vm["subgraph_0"](*self.tvm_input)
         self.counter += 1
-        return torch.from_numpy(output)
+        return torch.from_dlpack(output)
 
     def compile(self, module, **kwargs):
         with torch.no_grad():

graph_net/torch/test_compiler.py

@@ -251,7 +251,6 @@ def test_single_model(args):
 
     expected_out = eager_model_call()
     compiled_out = compiled_model_call()
-    compiled_out = (tensor.to(args.device) for tensor in compiled_out)
 
     def print_and_store_cmp(key, func, **kwargs):
         cmp_ret = func(expected_out, compiled_out, **kwargs)