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
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
.
1 Like