デバイス情報を取得できる `nn.Module`

module_with_device.py

import torch
import torch.nn as nn
import torch.backends.mps


class ModuleWithDevice(nn.Module):
    def __init__(self, *args, default_device=torch.device("cpu"), **kwargs) -> None:
        super().__init__(*args, **kwargs)
        # If this module has no parameters, returns this value.
        self._default_device = default_device

    @property
    def device(self) -> torch.device:
        for param in self.parameters():
            return param.device

        return self._default_device


class Encoder(ModuleWithDevice):
    def __init__(self) -> None:
        super().__init__()
        self.layer = nn.Linear(10, 20)


if __name__ == "__main__":
    encoder = Encoder()
    print(encoder.device)

    if torch.cuda.is_available():
        device = "cuda"
    elif torch.backends.mps.is_available():
        device = "mps"
    else:
        device = "cpu"

    encoder.to(device)

    print(encoder.device)