Add tests

Author: kaushikb11

tests/accelerators/test_accelerator_connector.py

@@ -41,8 +41,8 @@
     SLURMEnvironment,
     TorchElasticEnvironment,
 )
-from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities import _GPU_AVAILABLE
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.helpers.boring_model import BoringModel
 from tests.helpers.runif import RunIf
 
@@ -562,9 +562,9 @@ def test_accelerator_cpu():
 
     with pytest.raises(MisconfigurationException, match="You passed `accelerator='gpu'`, but GPUs are not available"):
         trainer = Trainer(accelerator="gpu")
-        
-    # with pytest.raises(MisconfigurationException, match="You passed `accelerator='gpu'`, but GPUs are not available"):
-    #     trainer = Trainer(accelerator="cpu", gpus=1)
+
+    with pytest.raises(MisconfigurationException, match="You requested GPUs:"):
+        trainer = Trainer(accelerator="cpu", gpus=1)
 
 
 @RunIf(min_gpus=1)
@@ -575,7 +575,9 @@ def test_accelerator_gpu():
     assert trainer._device_type == "gpu"
     assert isinstance(trainer.accelerator, GPUAccelerator)
 
-    with pytest.raises(MisconfigurationException, match="You passed `accelerator='gpu'`, but you didn't pass `gpus` to `Trainer`"):
+    with pytest.raises(
+        MisconfigurationException, match="You passed `accelerator='gpu'`, but you didn't pass `gpus` to `Trainer`"
+    ):
         trainer = Trainer(accelerator="gpu")
 
     trainer = Trainer(accelerator="auto", gpus=1)
@@ -584,9 +586,10 @@ def test_accelerator_gpu():
     assert isinstance(trainer.accelerator, GPUAccelerator)
 
 
+@RunIf(min_gpus=1)
 def test_accelerator_cpu_with_gpus_flag():
-    
+
     trainer = Trainer(accelerator="cpu", gpus=1)
 
     assert trainer._device_type == "cpu"
-    assert isinstance(trainer.accelerator, GPUAccelerator)
+    assert isinstance(trainer.accelerator, CPUAccelerator)

tests/accelerators/test_tpu_backend.py

@@ -16,6 +16,9 @@
 from torch import nn
 
 from pytorch_lightning import Trainer
+from pytorch_lightning.accelerators.cpu import CPUAccelerator
+from pytorch_lightning.accelerators.tpu import TPUAccelerator
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.helpers.boring_model import BoringModel
 from tests.helpers.runif import RunIf
 from tests.helpers.utils import pl_multi_process_test
@@ -113,3 +116,31 @@ def on_post_move_to_device(self):
 
     with pytest.warns(UserWarning, match="The model layers do not match"):
         trainer.fit(model)
+
+
+@RunIf(tpu=True)
+def test_accelerator_tpu():
+
+    trainer = Trainer(accelerator="tpu", tpu_cores=8)
+
+    assert trainer._device_type == "tpu"
+    assert isinstance(trainer.accelerator, TPUAccelerator)
+
+    with pytest.raises(
+        MisconfigurationException, match="You passed `accelerator='tpu'`, but you didn't pass `tpu_cores` to `Trainer`"
+    ):
+        trainer = Trainer(accelerator="tpu")
+
+    trainer = Trainer(accelerator="auto", tpu_cores=8)
+
+    assert trainer._device_type == "tpu"
+    assert isinstance(trainer.accelerator, TPUAccelerator)
+
+
+@RunIf(tpu=True)
+def test_accelerator_cpu_with_tpu_cores_flag():
+
+    trainer = Trainer(accelerator="cpu", tpu_cores=8)
+
+    assert trainer._device_type == "cpu"
+    assert isinstance(trainer.accelerator, CPUAccelerator)