From 1d05e8e0442894374faf47ed86cbe1037149f3d0 Mon Sep 17 00:00:00 2001 From: Amit Date: Sat, 8 Apr 2023 13:03:21 +0300 Subject: [PATCH] Update pytorch_builder.py Fix #100 --- hiddenlayer/pytorch_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/hiddenlayer/pytorch_builder.py b/hiddenlayer/pytorch_builder.py index 702c167..96c38cd 100644 --- 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: