Force hook standalone tests to single device

Author: lantiga

tests/tests_pytorch/models/test_hooks.py

@@ -61,7 +61,12 @@ def on_before_zero_grad(self, optimizer):
 
     model = CurrentTestModel()
 
-    trainer = Trainer(default_root_dir=tmp_path, max_steps=max_steps, max_epochs=2)
+    trainer = Trainer(
+        devices=1,
+        default_root_dir=tmp_path,
+        max_steps=max_steps,
+        max_epochs=2
+    )
     assert model.on_before_zero_grad_called == 0
     trainer.fit(model)
     assert max_steps == model.on_before_zero_grad_called
@@ -406,7 +411,7 @@ def prepare_data(self): ...
 @pytest.mark.parametrize(
     "kwargs",
     [
-        {},
+        {"devices": 1},
         # these precision plugins modify the optimization flow, so testing them explicitly
         pytest.param({"accelerator": "gpu", "devices": 1, "precision": "16-mixed"}, marks=RunIf(min_cuda_gpus=1)),
         pytest.param(
@@ -528,6 +533,7 @@ def test_trainer_model_hook_system_fit_no_val_and_resume_max_epochs(tmp_path):
     # initial training to get a checkpoint
     model = BoringModel()
     trainer = Trainer(
+        devices=1,
         default_root_dir=tmp_path,
         max_epochs=1,
         limit_train_batches=2,
@@ -543,6 +549,7 @@ def test_trainer_model_hook_system_fit_no_val_and_resume_max_epochs(tmp_path):
     callback = HookedCallback(called)
     # already performed 1 step, resume and do 2 more
     trainer = Trainer(
+        devices=1,
         default_root_dir=tmp_path,
         max_epochs=2,
         limit_train_batches=2,
@@ -605,6 +612,7 @@ def test_trainer_model_hook_system_fit_no_val_and_resume_max_steps(tmp_path):
     # initial training to get a checkpoint
     model = BoringModel()
     trainer = Trainer(
+        devices=1,
         default_root_dir=tmp_path,
         max_steps=1,
         limit_val_batches=0,
@@ -624,6 +632,7 @@ def test_trainer_model_hook_system_fit_no_val_and_resume_max_steps(tmp_path):
     train_batches = 2
     steps_after_reload = 1 + train_batches
     trainer = Trainer(
+        devices=1,
         default_root_dir=tmp_path,
         max_steps=steps_after_reload,
         limit_val_batches=0,
@@ -690,6 +699,7 @@ def test_trainer_model_hook_system_eval(tmp_path, override_on_x_model_train, bat
     assert is_overridden(f"on_{noun}_model_train", model) == override_on_x_model_train
     callback = HookedCallback(called)
     trainer = Trainer(
+        devices=1,
         default_root_dir=tmp_path,
         max_epochs=1,
         limit_val_batches=batches,
@@ -731,7 +741,10 @@ def test_trainer_model_hook_system_predict(tmp_path):
     callback = HookedCallback(called)
     batches = 2
     trainer = Trainer(
-        default_root_dir=tmp_path, limit_predict_batches=batches, enable_progress_bar=False, callbacks=[callback]
+        devices=1,
+        default_root_dir=tmp_path,
+        limit_predict_batches=batches,
+        enable_progress_bar=False, callbacks=[callback]
     )
     trainer.predict(model)
     expected = [
@@ -797,7 +810,7 @@ def predict_dataloader(self):
 
     model = CustomBoringModel()
 
-    trainer = Trainer(default_root_dir=tmp_path, fast_dev_run=5)
+    trainer = Trainer(devices=1, default_root_dir=tmp_path, fast_dev_run=5)
 
     trainer.fit(model)
     trainer.test(model)
@@ -812,6 +825,7 @@ def test_trainer_datamodule_hook_system(tmp_path):
     model = BoringModel()
     batches = 2
     trainer = Trainer(
+        devices=1,
         default_root_dir=tmp_path,
         max_epochs=1,
         limit_train_batches=batches,
@@ -887,7 +901,7 @@ class CustomHookedModel(HookedModel):
     assert is_overridden("configure_model", model) == override_configure_model
 
     datamodule = CustomHookedDataModule(ldm_called)
-    trainer = Trainer()
+    trainer = Trainer(devices=1)
     trainer.strategy.connect(model)
     trainer._data_connector.attach_data(model, datamodule=datamodule)
     ckpt_path = str(tmp_path / "file.ckpt")
@@ -960,6 +974,7 @@ def predict_step(self, *args, **kwargs):
 
     model = MixedTrainModeModule()
     trainer = Trainer(
+        devices=1,
         default_root_dir=tmp_path,
         max_epochs=1,
         val_check_interval=1,