chore: Update torch2onnx.py

yakhyo · May 19, 2024 · f2ece5e · f2ece5e
1 parent c5948b5
commit f2ece5e
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/utils/torch2onnx.py b/utils/torch2onnx.py
@@ -10,6 +10,7 @@ def torch2onnx_export(params):
 
     model = BiSeNet(num_classes, backbone_name=params.model)
     model.load_state_dict(torch.load(params.weight))
+    model.eval()
 
     onnx_model_path = params.weight.replace(".pt", ".onnx")