add torch 2.2.2 or below compatibility

Author: GdoongMathew

src/lightning/fabric/utilities/device_dtype_mixin.py

@@ -25,7 +25,9 @@ class _DeviceDtypeModuleMixin(Module):
     def __init__(self) -> None:
         super().__init__()
         self._dtype: Union[str, torch.dtype] = torch.get_default_dtype()
-        self._device: torch.device = torch.get_default_device()
+        # Workarounds from the original pytorch issue:
+        # https://github.com/pytorch/pytorch/issues/115333#issuecomment-1848449687
+        self._device = torch.get_default_device() if torch.__version__ >= "2.3.0" else torch.empty(0).device
 
     @property
     def dtype(self) -> Union[str, torch.dtype]:

tests/tests_fabric/utilities/test_device_dtype_mixin.py

@@ -50,6 +50,27 @@ def test_submodules_device_and_dtype(dst_device_str, dst_type):
     assert model.dtype == model.module.module.dtype == dst_type
 
 
+@pytest.mark.parametrize(
+    "dst_device_str",
+    [
+        "cpu",
+        pytest.param("cuda:0", marks=RunIf(min_cuda_gpus=1)),
+        pytest.param("mps:0", marks=RunIf(mps=True)),
+    ],
+)
+@pytest.mark.parametrize("dst_type", [torch.half, torch.float, torch.double])
+def test_submodules_context_device_and_dtype(dst_device_str, dst_type):
+    if dst_device_str == "mps:0" and dst_type in (torch.half, torch.double):
+        pytest.skip("MPS does not yet support half and double.")
+
+    dst_device = torch.device(dst_device_str)
+    torch.set_default_dtype(dst_type)
+    with dst_device:
+        model = TopModule()
+    assert model.device == dst_device
+    assert model.dtype == dst_type
+
+
 @pytest.mark.parametrize(
     "device",
     [