Fix tests

sillsdev · Feb 3, 2024 · fd774d6 · fd774d6
1 parent 850bca4
commit fd774d6
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/machine/jobs/huggingface/hugging_face_nmt_model_factory.py b/machine/jobs/huggingface/hugging_face_nmt_model_factory.py
@@ -85,7 +85,7 @@ def create_engine(self) -> TranslationEngine:
         )
 
     def save_model(self) -> None:
-        if self._config.save_model is None:
+        if "save_model" not in self._config:
             return
 
         tar_file_path = Path(self._config.data_dir, "builds", self._config.build_id, "model.tar.gz")

diff --git a/machine/jobs/nmt_engine_build_job.py b/machine/jobs/nmt_engine_build_job.py
@@ -94,7 +94,7 @@ def run(
                 current_inference_step += len(pi_batch)
                 phase_progress(ProgressStatus.from_step(current_inference_step, inference_step_count))
 
-        if self._config.save_model is not None:
+        if "save_model" in self._config and self._config.save_model is not None:
             logger.info("Saving model")
             self._nmt_model_factory.save_model()