Repository: Lightning-AI/pytorch-lightning
File: src/lightning/fabric/utilities/device_dtype_mixin.py (Python)

Code before:
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, Optional, Union

import torch
from torch.nn import Module
from typing_extensions import Self, override


class _DeviceDtypeModuleMixin(Module):
    __jit_unused_properties__: list[str] = ["device", "dtype"]

    def __init__(self) -> None:
        super().__init__()
        self._dtype: Union[str, torch.dtype] = torch.get_default_dtype()
        # 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]:
        return self._dtype

    @dtype.setter
    def dtype(self, new_dtype: Union[str, torch.dtype]) -> None:
        # necessary to avoid infinite recursion
        raise RuntimeError("Cannot set the dtype explicitly. Please use module.to(new_dtype).")

    @property
    def device(self) -> torch.device:
        device = self._device

        # make this more explicit to always include the index
        if device.type == "cuda" and device.index is None:
            return torch.device(f"cuda:{torch.cuda.current_device()}")

        return device

    @override
    def to(self, *args: Any, **kwargs: Any) -> Self:
        """See :meth:`torch.nn.Module.to`."""
        # this converts `str` device to `torch.device`
        device, dtype = torch._C._nn._parse_to(*args, **kwargs)[:2]

Review comment (suggestion):
We already seem to have the `_TORCH_GREATER_EQUAL_2_3` constant defined in
`lightning.fabric.utilities.imports`, so we can reuse it here.

https://github.com/Lightning-AI/pytorch-lightning/blob/1fc077b6e4f63466cea3fad33dda1359920b7287/src/lightning/fabric/utilities/imports.py#L32-L34

---

```suggestion
        self._device = torch.get_default_device() if _TORCH_GREATER_EQUAL_2_3 else torch.empty(0).device
```

```python 
# import 
from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_3
```


Code after:
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, Optional, Union

import torch
from torch.nn import Module
from typing_extensions import Self, override

from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_3


class _DeviceDtypeModuleMixin(Module):
    __jit_unused_properties__: list[str] = ["device", "dtype"]

    def __init__(self) -> None:
        super().__init__()
        self._dtype: Union[str, torch.dtype] = torch.get_default_dtype()
        # Workarounds from the original pytorch issue:
        # https://github.com/pytorch/pytorch/issues/115333#issuecomment-1848449687
        self._device = torch.get_default_device() if _TORCH_GREATER_EQUAL_2_3 else torch.empty(0).device

    @property
    def dtype(self) -> Union[str, torch.dtype]:
        return self._dtype

    @dtype.setter
    def dtype(self, new_dtype: Union[str, torch.dtype]) -> None:
        # necessary to avoid infinite recursion
        raise RuntimeError("Cannot set the dtype explicitly. Please use module.to(new_dtype).")

    @property
    def device(self) -> torch.device:
        device = self._device

        # make this more explicit to always include the index
        if device.type == "cuda" and device.index is None:
            return torch.device(f"cuda:{torch.cuda.current_device()}")

        return device

    @override
    def to(self, *args: Any, **kwargs: Any) -> Self:
        """See :meth:`torch.nn.Module.to`."""