clean

mosaicml · Dec 1, 2024 · 6785988 · 6785988
1 parent ceb6591
commit 6785988
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_events.py b/tests/test_events.py
@@ -25,7 +25,7 @@ class TestEventCalls:
     eval_subset_num_batches = 1
     train_subset_num_batches = 1
 
-    def get_trainer(self, precision='fp32', max_duration='1ep', **kwargs):
+    def get_trainer(self, precision='fp32', **kwargs):
         model = SimpleModel()
         optimizer = torch.optim.Adam(model.parameters())
 
@@ -57,7 +57,7 @@ def get_trainer(self, precision='fp32', max_duration='1ep', **kwargs):
             precision=precision,
             train_subset_num_batches=self.train_subset_num_batches,
             eval_subset_num_batches=self.eval_subset_num_batches,
-            max_duration=max_duration,
+            max_duration='1ep',
             optimizers=optimizer,
             callbacks=[EventCounterCallback()],
             **kwargs,