Skip to content

Commit ceca5ef

Browse files
committed
Remove torch_out from onnx export, no point without the export_ fn
1 parent 0cae8a4 commit ceca5ef

File tree

1 file changed

+2
-7
lines changed

1 file changed

+2
-7
lines changed

timm/utils/onnx.py

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -78,9 +78,8 @@ def onnx_export(
7878
export_options=export_options,
7979
)
8080
export_output.save(output_file)
81-
torch_out = None
8281
else:
83-
torch_out = torch.onnx.export(
82+
torch.onnx.export(
8483
model,
8584
example_input,
8685
output_file,
@@ -101,9 +100,5 @@ def onnx_export(
101100
if check_forward and not training:
102101
import numpy as np
103102
onnx_out = onnx_forward(output_file, example_input)
104-
if torch_out is not None:
105-
np.testing.assert_almost_equal(torch_out.numpy(), onnx_out, decimal=3)
106-
np.testing.assert_almost_equal(original_out.numpy(), torch_out.numpy(), decimal=5)
107-
else:
108-
np.testing.assert_almost_equal(original_out.numpy(), onnx_out, decimal=3)
103+
np.testing.assert_almost_equal(original_out.numpy(), onnx_out, decimal=3)
109104

0 commit comments

Comments
 (0)