feat: Add save_strategy

alexandrainst · Nov 23, 2023 · ef5d4d5 · ef5d4d5
1 parent ce7780b
commit ef5d4d5
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/coral_models/wav2vec2.py b/src/coral_models/wav2vec2.py
@@ -237,6 +237,7 @@ def load_training_arguments(self) -> TrainingArguments:
             evaluation_strategy="steps" if do_eval else "no",
             eval_steps=self.cfg.eval_steps if do_eval else None,
             save_steps=self.cfg.save_steps,
+            save_strategy="no" if self.cfg.save_total_limit == 0 else "steps",
             logging_steps=self.cfg.logging_steps,
             length_column_name="input_length",
             gradient_checkpointing=True,