diffusion_policy_gbc / model /common /module_attr_mixin.py
ducido's picture
Initial commit
987ed1b verified
Raw
History Blame Contribute Delete
334 Bytes
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