Update pytorch_builder.py

Fix waleedka#100
amitbar05 · Apr 8, 2023 · 1d05e8e · 1d05e8e
1 parent 45243d5
commit 1d05e8e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/hiddenlayer/pytorch_builder.py b/hiddenlayer/pytorch_builder.py
@@ -68,7 +68,7 @@ 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)
+    torch_graph = torch.onnx._optimize_graph(trace, torch.onnx.OperatorExportTypes.ONNX)
 
     # Dump list of nodes (DEBUG only)
     if verbose: