[TPU] Fix test assertion error from artifacts (#19825)

Lightning-AI · May 23, 2024 · 7874cd0 · 7874cd0
1 parent e0d7ede
commit 7874cd0
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/tests/tests_pytorch/accelerators/test_xla.py b/tests/tests_pytorch/accelerators/test_xla.py
@@ -56,7 +56,7 @@ def test_resume_training_on_cpu(tmp_path):
     """Checks if training can be resumed from a saved checkpoint on CPU."""
     # Train a model on TPU
     model = BoringModel()
-    trainer = Trainer(max_epochs=1, accelerator="tpu", devices="auto")
+    trainer = Trainer(max_epochs=1, accelerator="tpu", devices="auto", default_root_dir=tmp_path)
     trainer.fit(model)
 
     if trainer.world_size != trainer.num_devices:

diff --git a/tests/tests_pytorch/trainer/properties/test_estimated_stepping_batches.py b/tests/tests_pytorch/trainer/properties/test_estimated_stepping_batches.py
@@ -152,6 +152,6 @@ def on_train_start(self):
 @mock.patch.dict(os.environ, os.environ.copy(), clear=True)
 def test_num_stepping_batches_with_tpu_multi():
     """Test stepping batches with the TPU strategy across multiple devices."""
-    trainer = Trainer(accelerator="tpu", devices="auto", max_epochs=1)
+    trainer = Trainer(accelerator="tpu", devices="auto", max_epochs=1, logger=False, enable_checkpointing=False)
     model = MultiprocessModel()
     trainer.fit(model)