File size: 334 Bytes
987ed1b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 | import torch.nn as nn
class ModuleAttrMixin(nn.Module):
def __init__(self):
super().__init__()
self._dummy_variable = nn.Parameter()
@property
def device(self):
return next(iter(self.parameters())).device
@property
def dtype(self):
return next(iter(self.parameters())).dtype
|