Save_ONNX defaults to False

AdvancedPhotonSource · Apr 26, 2024 · 75be2ae · 75be2ae
1 parent adf8753
commit 75be2ae
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/generic_trainer/configs.py b/generic_trainer/configs.py
@@ -252,7 +252,7 @@ class TrainingConfig(Config):
     automatic_mixed_precision: bool = False
     """Automatic mixed precision and gradient scaling are enabled if True."""
 
-    save_onnx: bool = True
+    save_onnx: bool = False
     """If True, ONNX models are saved along with state dicts."""
 
 

diff --git a/generic_trainer/trainer.py b/generic_trainer/trainer.py
@@ -836,7 +836,7 @@ def save_model(self, path, subcomponent=None):
             m = getattr(m, subcomponent)
         torch.save(m.state_dict(), path)
 
-    def update_saved_model(self, filename='best_model.pth', save_configs=True, save_onnx=True, subcomponent=None,
+    def update_saved_model(self, filename='best_model.pth', save_configs=True, save_onnx=False, subcomponent=None,
                            run_with_only_rank_0=True):
         """
         Updates saved model if validation loss is minimum.