Fix TypeError for Pytorch Model waleedka#94

julien-blanchon · Jun 23, 2022 · 8e14b20 · 8e14b20
1 parent 45243d5
commit 8e14b20
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/hiddenlayer/pytorch_builder.py b/hiddenlayer/pytorch_builder.py
@@ -68,7 +68,10 @@ def import_graph(hl_graph, model, args, input_names=None, verbose=False):
 
     # Run the Pytorch graph to get a trace and generate a graph from it
     trace, out = torch.jit._get_trace_graph(model, args)
-    torch_graph = torch.onnx._optimize_trace(trace, torch.onnx.OperatorExportTypes.ONNX)
+    try:
+        torch_graph = torch.onnx._optimize_trace(trace, torch.onnx.OperatorExportTypes.ONNX)
+    except TypeError as e:
+        torch_graph = trace    
 
     # Dump list of nodes (DEBUG only)
     if verbose: