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.
EleniXoch · May 6, 2024 · ca41b2b · ca41b2b
1 parent 566d07f
commit ca41b2b
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"],