update

Author: Borda

tests/tests_pytorch/callbacks/test_prediction_writer.py

@@ -83,7 +83,7 @@ def test_prediction_writer_batch_indices(num_workers, tmp_path):
     DummyPredictionWriter.write_on_batch_end = Mock()
     DummyPredictionWriter.write_on_epoch_end = Mock()
 
-    dataloader = DataLoader(RandomDataset(32, 64), batch_size=4, num_workers=num_workers, persistent_workers=True)
+    dataloader = DataLoader(RandomDataset(32, 64), batch_size=4, num_workers=num_workers, persistent_workers=num_workers > 0)
     model = BoringModel()
     writer = DummyPredictionWriter("batch_and_epoch")
     trainer = Trainer(default_root_dir=tmp_path, logger=False, limit_predict_batches=4, callbacks=writer)

tests/tests_pytorch/trainer/connectors/test_data_connector.py

@@ -135,7 +135,7 @@ def test_dataloader_persistent_workers_performance_warning(num_workers, tmp_path
         barebones=True,
     )
     model = TestSpawnBoringModel(warning_expected=(num_workers > 0))
-    dataloader = DataLoader(RandomDataset(32, 64), num_workers=num_workers, persistent_workers=True)
+    dataloader = DataLoader(RandomDataset(32, 64), num_workers=num_workers, persistent_workers=num_workers > 0)
     trainer.fit(model, dataloader)
 
 

tests/tests_pytorch/trainer/test_dataloaders.py

@@ -658,10 +658,7 @@ def on_train_epoch_end(self):
 def test_auto_add_worker_init_fn_distributed(tmp_path, monkeypatch):
     """Test that the lightning worker_init_fn takes care of dataloaders in multi-gpu/multi-node training."""
     dataset = NumpyRandomDataset()
-    num_workers = 2
-    batch_size = 2
-
-    dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers, persistent_workers=True)
+    dataloader = DataLoader(dataset, batch_size=2, num_workers=2, persistent_workers=True)
     seed_everything(0, workers=True)
     trainer = Trainer(default_root_dir=tmp_path, max_epochs=1, accelerator="gpu", devices=2, strategy="ddp_spawn")
     model = MultiProcessModel()