Bases: Module
This module is a base class for all modules that need to be compatible with this project.
It ensures that the module stores its current device and dtype.
Source code in bfm/model/base.py
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35 | class BFModule(nn.Module):
"""
This module is a base class for all modules that need to be compatible with this project.
It ensures that the module stores its current device and dtype.
"""
def __init__(self):
super().__init__()
self._device = None
self._dtype = None
def to(self, *args, **kwargs):
output = super().to(*args, **kwargs)
# Extract device and dtype from args/kwargs
device = next((torch.device(arg) for arg in args if isinstance(arg, (torch.device, str))),
kwargs.get('device', None))
dtype = next((arg for arg in args if isinstance(arg, torch.dtype)),
kwargs.get('dtype', None))
if device is not None: self._device = device
if dtype is not None: self._dtype = dtype
return output
@property
def device(self):
if self._device is None:
self._device = next(self.parameters()).device
return self._device
@property
def dtype(self):
if self._dtype is None:
self._dtype = next(self.parameters()).dtype
return self._dtype
|