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.
Yes all of the factory functions like randn or ones are independent of the nn.Module device.
So in general, you want to provide a device= argument to these functions if you want the Tensor to be created on a special device.
You can get the right device from any param or input of your Module.
Also, if you want Tensor attributes to be moved along with the parameters, you should register them as “buffers” on the Module with self.register_buffer("q_transposed", the_tensor) which you can then access as self.q_transposed.
Side not in your code sample, there is not need to wrap the Tensor into a Parameter here. Parameter are only important when you assign them to a nn.Module.
