precommit

mosaicml · Dec 1, 2024 · ceb6591 · ceb6591
1 parent 868a71b
commit ceb6591
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_events.py b/tests/test_events.py
@@ -31,7 +31,7 @@ def get_trainer(self, precision='fp32', max_duration='1ep', **kwargs):
 
         train_dataset = RandomClassificationDataset()
         eval_dataset = RandomClassificationDataset()
-        train_batch_size = 2 
+        train_batch_size = 2
 
         evaluator1 = DataLoader(
             dataset=eval_dataset,
@@ -99,7 +99,7 @@ def get_trainer(self, precision='fp32', max_duration='1ep', **kwargs):
             ),
         ],
     )
-    @pytest.mark.parametrize('save_interval', ['1ep'])
+    @pytest.mark.parametrize('save_interval', ['1ep', '1ba'])
     def test_event_calls(self, world_size, device, deepspeed_zero_stage, use_fsdp, precision, save_interval):
         save_interval = Time.from_timestring(save_interval)