update num epochs

mosaicml · Nov 30, 2024 · 69f13ba · 69f13ba
1 parent 62f0bea
commit 69f13ba
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_events.py b/tests/test_events.py
@@ -57,7 +57,7 @@ def get_trainer(self, precision='fp32', **kwargs):
             precision=precision,
             train_subset_num_batches=self.train_subset_num_batches,
             eval_subset_num_batches=self.eval_subset_num_batches,
-            max_duration='2ep',
+            max_duration='1ep',
             optimizers=optimizer,
             callbacks=[EventCounterCallback()],
             **kwargs,
@@ -127,7 +127,7 @@ def test_event_calls(self, world_size, device, deepspeed_zero_stage, use_fsdp, p
         )
         trainer.fit()
 
-        self._assert_expected_event_calls(trainer, save_interval, num_epochs=2)
+        self._assert_expected_event_calls(trainer, save_interval, num_epochs=1)
 
     def _assert_expected_event_calls(self, trainer: Trainer, eval_interval: Time, num_epochs: int):
         state = trainer.state