Random tensor inside a cuda module doesn't get created on cuda


#1

I have a model (nn.module) class and inside this class i create a random tensor using torch.randn(). When I create a model.to(device) where device is cuda, the random tensor doesn’t get moved to cuda. Is the behavior correct? how should i do it, without passing device to the model.
Thanks


#2

By default only Parameters will be recognized and pushed to the according device.
Try to wrap your random tensor into torch.nn.Parameter(torch.randn(1), requires_grad=False) and it should work.
Alternatively you could also store it in a nn.ModuleList.