I was implementing a torch.qr inside a module of my creation with something like this:
self.q_transposed=torch.qr(nn.Parameter(torch.randn(self.m,self.n),requires_grad=False).T)[0].T
Even though the module device is cuda (I made sure of that), the parameter is in cpu for some reason.
The same thing is happening with torch.ones(): inside my forward method, I need to declare a torch.ones() tensor and then multiply it by another tensor (that is located in gpu) and, for some reason, this “ones” tensor goes directly to cpu.
The only solution I have found so far was to insert a device argument in the init and then add each parameter back to cuda when my matmul would not work. But I am afraid that this would compromise using resources such as nn.DataParallel, that are going to be useful for my application.