fix: Wrong network size written to onnx (acts-project#3072)

This PR fix an issue for the seed selection NN where the wrong network size would be written to the onnx file. The default network are unaffected as they where written before this mistake was introduced.
Ragansu · Apr 19, 2024 · f725096 · f725096
1 parent d452d17
commit f725096
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"],