Fix daily test for iteration (#3068)

mosaicml · Feb 28, 2024 · 85a4bc1 · 85a4bc1
1 parent 6c833c6
commit 85a4bc1
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tests/trainer/test_checkpoint.py b/tests/trainer/test_checkpoint.py
@@ -864,6 +864,9 @@ def test_load_remote_checkpoint(self, device, tmp_path: pathlib.Path, load_weigh
             device=device,
         )
 
+        # TODO(GRT-2735): Update remote checkpoint with iteration.
+        trainer_2.state.timestamp._epoch_in_iteration = Time.from_input(2, TimeUnit.EPOCH)
+
         # check weights loaded properly
         self._assert_weights_equivalent(
             trainer_1.state.model,