From b8bc8745ae6c025255cec490f27a42d696ed2de4 Mon Sep 17 00:00:00 2001
From: Yaman Umuroglu <yamanu@xilinx.com>
Date: Tue, 3 Sep 2019 23:26:57 +0100
Subject: [PATCH] put network into eval mode before exporting
---
tests/test_brevitas_export.py | 1 +
1 file changed, 1 insertion(+)
diff --git a/tests/test_brevitas_export.py b/tests/test_brevitas_export.py
index 23d8af3d1..b3aa98793 100644
--- a/tests/test_brevitas_export.py
+++ b/tests/test_brevitas_export.py
@@ -74,6 +74,7 @@ def test_brevitas_to_onnx_export():
export_onnx_path = "test_output_lfc.onnx"
with torch.no_grad():
lfc = LFC(weight_bit_width=1, act_bit_width=1, in_bit_width=1)
+ lfc = lfc.eval()
bo.prepare_for_onnx_export(lfc, True)
torch.onnx.export(
lfc, torch.empty(784, dtype=torch.float), export_onnx_path, verbose=True
--
GitLab