fix

Xilinx · Oct 31, 2024 · 9bf71fe · 9bf71fe
1 parent cba513c
commit 9bf71fe
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/brevitas/export/onnx/manager.py b/src/brevitas/export/onnx/manager.py
@@ -171,7 +171,7 @@ def export_onnx(
 
                     with PatchFp8Ops():
                         torch.onnx.export(module, args, export_target, **onnx_export_kwargs)
-                    restore_return_quant_tensor(model, return_quant_tensor_state)
+                    restore_return_quant_tensor(module, return_quant_tensor_state)
                     # restore the model to previous properties
                     module.apply(lambda m: _restore_act_caching_mode(m))
                     cls.set_export_mode(module, enabled=False)