fix tests.

Author: GdoongMathew

tests/tests_pytorch/accelerators/test_gpu.py

@@ -68,3 +68,17 @@ def test_gpu_availability():
 def test_warning_if_gpus_not_used(cuda_count_1):
     with pytest.warns(UserWarning, match="GPU available but not used"):
         Trainer(accelerator="cpu")
+
+
+@RunIf(min_cuda_gpus=1)
+def test_gpu_device_name():
+    for i in range(torch.cuda.device_count()):
+        assert torch.cuda.get_device_name(i) == CUDAAccelerator.device_name(i)
+
+    with torch.device("cuda:0"):
+        assert torch.cuda.get_device_name(0) == CUDAAccelerator.device_name()
+
+
+def test_gpu_device_name_no_gpu():
+    with mock.patch("torch.cuda.is_available", return_value=False):
+        assert str(False) == CUDAAccelerator.device_name()

tests/tests_pytorch/accelerators/test_mps.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from collections import namedtuple
+from unittest import mock
 
 import pytest
 import torch
@@ -39,6 +40,16 @@ def test_mps_availability():
     assert MPSAccelerator.is_available()
 
 
+@RunIf(mps=True)
+def test_mps_device_name():
+    assert MPSAccelerator.device_name() == "True (mps)"
+
+
+def test_mps_device_name_not_available():
+    with mock.patch("torch.backends.mps.is_available", return_value=False):
+        assert MPSAccelerator.device_name() == "False"
+
+
 def test_warning_if_mps_not_used(mps_count_1):
     with pytest.warns(UserWarning, match="GPU available but not used"):
         Trainer(accelerator="cpu")

tests/tests_pytorch/accelerators/test_xla.py

@@ -302,6 +302,11 @@ def test_warning_if_tpus_not_used(tpu_available):
         Trainer(accelerator="cpu")
 
 
+@RunIf(tpu=True)
+def test_tpu_device_name():
+    assert XLAAccelerator.device_name() == "TPU"
+
+
 @pytest.mark.parametrize(
     ("devices", "expected_device_ids"),
     [

tests/tests_pytorch/conftest.py

@@ -182,6 +182,7 @@ def thread_police_duuu_daaa_duuu_daaa():
 def mock_cuda_count(monkeypatch, n: int) -> None:
     monkeypatch.setattr(lightning.fabric.accelerators.cuda, "num_cuda_devices", lambda: n)
     monkeypatch.setattr(lightning.pytorch.accelerators.cuda, "num_cuda_devices", lambda: n)
+    monkeypatch.setattr(torch.cuda, "get_device_name", lambda _: "Mocked CUDA Device")
 
 
 @pytest.fixture
@@ -244,6 +245,7 @@ def mock_tpu_available(monkeypatch: pytest.MonkeyPatch, value: bool = True) -> N
     monkeypatch.setitem(sys.modules, "torch_xla", Mock())
     monkeypatch.setitem(sys.modules, "torch_xla.core.xla_model", Mock())
     monkeypatch.setitem(sys.modules, "torch_xla.experimental", Mock())
+    monkeypatch.setattr(lightning.pytorch.accelerators.xla.XLAAccelerator, "device_name", lambda _: "TPU")
 
 
 @pytest.fixture