Skip to content

DeviceMixin

Methods and Attributes

Consistent device management across all modules.

This should be inherited by all classes.

Source code in torchebm/core/device_mixin.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
class DeviceMixin:
    """Consistent device management across all modules.

    This should be inherited by all classes.
    """

    def __init__(self, device: Union[str, torch.device, None] = None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._device = normalize_device(device)

    @property
    def device(self) -> torch.device:
        if self._device is not None:
            return normalize_device(self._device)
        if self._device is None:
            if hasattr(self, "parameters") and callable(getattr(self, "parameters")):
                try:
                    param_device = next(self.parameters()).device
                    return normalize_device(param_device)
                except StopIteration:
                    pass

            if hasattr(self, "buffers") and callable(getattr(self, "buffers")):
                try:
                    buffer_device = next(self.buffers()).device
                    return normalize_device(buffer_device)
                except StopIteration:
                    pass

        return torch.device("cuda" if torch.cuda.is_available() else "cpu")

    def to(self, *args, **kwargs):
        """Override to() to update internal device tracking."""
        result = super().to(*args, **kwargs)

        if args:
            if isinstance(args[0], (str, torch.device)):
                self._device = normalize_device(args[0])
        if "device" in kwargs:
            self._device = normalize_device(kwargs["device"])

        return result

_device instance-attribute 

_device = normalize_device(device)

device property 

device: device

to 

to(*args, **kwargs)

Override to() to update internal device tracking.

Source code in torchebm/core/device_mixin.py 
66
67
68
69
70
71
72
73
74
75
76
def to(self, *args, **kwargs):
    """Override to() to update internal device tracking."""
    result = super().to(*args, **kwargs)

    if args:
        if isinstance(args[0], (str, torch.device)):
            self._device = normalize_device(args[0])
    if "device" in kwargs:
        self._device = normalize_device(kwargs["device"])

    return result