From fbd8ba3ebc34bcd7ffc33fb882c52bd4c407a15b Mon Sep 17 00:00:00 2001
From: NewLandTV <86148752+NewLandTV@users.noreply.github.com>
Date: Tue, 20 Feb 2024 21:13:14 +0900
Subject: [PATCH] Fix ONNX Builder

---
 Build.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/Build.py b/Build.py
index f203aea..bcf2ee5 100644
--- a/Build.py
+++ b/Build.py
@@ -12,8 +12,9 @@
 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(
@@ -21,7 +22,7 @@
     x,
     path,
     export_params = True,
-    opset_version = 10,
+    opset_version = 13,
     do_constant_folding = True,
     input_names = ["input"],
     output_names = ["output"],