Torch FX graph outputs fix

cavusmustafa · Mar 31, 2023 · e00427a · e00427a
1 parent 035dc66
commit e00427a
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/backend.py b/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/backend.py
@@ -72,7 +72,7 @@ def fx_openvino(subgraph, example_inputs):
         compiled_model = partitioner.make_partitions(model)
 
         def _call(*args):
-            res = execute(compiled_model, *example_inputs, executor="openvino")
+            res = execute(compiled_model, *args, executor="openvino")
             return res
         return _call
     except Exception as e:

diff --git a/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/execute.py b/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/execute.py
@@ -72,8 +72,14 @@ def openvino_execute(gm: GraphModule, *args, executor_parameters=None, partition
     res = compiled(ov_inputs)
 
     results1 = [res[out] for out in compiled.outputs]
-    results = torch.from_numpy(np.array(results1, dtype=np.float32))
-    return results
+
+    if len(results1) == 1:
+        results = torch.from_numpy(results1[0])
+        return results
+    else:
+        for i in range(len(results1)):
+            results1[i] = torch.from_numpy(results1[i])
+        return results1
 
 
 class OpenVINOGraphModule(torch.nn.Module):