fix network size issue

acts-project · Apr 3, 2024 · a5aa62a · a5aa62a
1 parent 4271243
commit a5aa62a
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/Examples/Scripts/Python/MLAmbiguityResolution/train_ambiguity_solver.py b/Examples/Scripts/Python/MLAmbiguityResolution/train_ambiguity_solver.py
@@ -273,7 +273,7 @@ def train(
 torch.save(duplicateClassifier, "duplicateClassifier.pt")
 torch.onnx.export(
     duplicateClassifier,
-    input_test,
+    input_test[0:1],
     "duplicateClassifier.onnx",
     input_names=["x"],
     output_names=["y"],

diff --git a/Examples/Scripts/Python/MLAmbiguityResolution/train_seed_solver.py b/Examples/Scripts/Python/MLAmbiguityResolution/train_seed_solver.py
@@ -322,7 +322,7 @@ def train(
 torch.save(duplicateClassifier, "seedduplicateClassifier.pt")
 torch.onnx.export(
     duplicateClassifier,
-    input_test[0],
+    input_test[0:1],
     "seedduplicateClassifier.onnx",
     input_names=["x"],
     output_names=["y"],