update num epochs

Author: v-chen_data

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