You could derive the .cuda()  and/or the .to() function for something like this:
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


EDIT: I recommend only to derive the .to() function because you have to make sure that you also derive the .cpu() function when deriving .cuda()