gistfile1.txt

class Net(torch.nn.Module):
    def __init__(self):
        self.a = torch.nn.Parameter(torch.zeros(1))
        self.b = torch.zeros(1)

    def forward(self, inputs):
        #do some stuff
    
    def cuda(self, device=None):
        self = super().cuda(device)
        self.b = self.b.cuda(device)
        return self 

    def to(self, *args, **kwargs):
        self = super().to(*args, **kwargs) 
        self.b = self.b.to(*args, **kwargs) 
        return self