I have a model and I have put the model into cuda depending on the availability of gpu. If gpu is available, I put the model and the inputs x into cuda but the tensors in class A which are part of the model are getting loaded into cpu. How to put all the tensors used in class A into cuda if gpu is available? Please help!

class A(nn.Module):
def __init__(self, dim):

    super().__init__()

    self.dim = dim



    Q = torch.nn.init.orthogonal_(torch.randn(dim, dim))

    P, L, U = torch.lu_unpack(*Q.lu())

    self.P = P # remains fixed during optimization

    self.L = nn.Parameter(L) # lower triangular portion

    self.S = nn.Parameter(U.diag()) # "crop out" the diagonal to its own parameter

    self.U = nn.Parameter(torch.triu(U, diagonal=1)) # "crop out" diagonal, stored in S

def _assemble_W(self):

    """ assemble W from its pieces (P, L, U, S) """

    L = torch.tril(self.L, diagonal=-1) + torch.diag(torch.ones(self.dim))

    U = torch.triu(self.U, diagonal=1)

    W = self.P @ L @ (U + torch.diag(self.S))

    return W

def forward(self, x):

    W = self._assemble_W()

    z = x @ W

    log_det = torch.sum(torch.log(torch.abs(self.S)))

    return z, log_det