Fix ONNX Builder

NewLandTV · Feb 20, 2024 · fbd8ba3 · fbd8ba3
1 parent b091728
commit fbd8ba3
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/Build.py b/Build.py
@@ -12,16 +12,17 @@
 model.eval()
 
 # Export onnx
-x = torch.randn(batchSize, 1, 28, 28, requires_grad = True)
-out = model(x)
+x = torch.randn(batchSize, 3, 28, 28, requires_grad = True)
+outputs = model(x)
+_, predicted = torch.max(outputs.data, 1)
 path = "Model.onnx"
 
 torch.onnx.export(
     model,
     x,
     path,
     export_params = True,
-    opset_version = 10,
+    opset_version = 13,
     do_constant_folding = True,
     input_names = ["input"],
     output_names = ["output"],