CODEWITHSUNDEEP
pythonpytorchgpu
zzzjjh
zzzjjh

Reputation: 21

pytorch .cuda() can't get the tensor to cuda

I try to get my data onto gpu,but it doesn't work;

in my train.py

if __name__ == '__main__':
ten = torch.FloatTensor(2)
ten = ten.cuda()
print(ten)

args = config()
train_net(args, args.train_net, loss_config=net_loss_config[args.train_net])

when it runs ,it prints this

tensor([0., 0.])

the tensor is not on cuda

but in test.py

import torch
ten=torch.FloatTensor(2)
ten=ten.cuda()
print(ten)

it prints this

tensor([1.4013e-45, 0.0000e+00], device='cuda:0')

now the tensor is on cuda

Upvotes: 2

Views: 1833

Answers (1)

Yashi Aggarwal
Yashi Aggarwal

Reputation: 405

The error means that the ten variable in your model is of type torch.FloatTensor (CPU), while the input you provide to the model is of type torch.cuda.FloatTensor (GPU).

The most likely scenario is that you have nn.Parameter or other modules such as nn.Conv2d defined in the __init__() method of your model, and additional weights or layers defined in the forward() method of your model.

In this case, the layers defined in the forward() method are not modules of the model, and they won’t be mapped to GPU when you call cuda().

You also need to explicitly add the parameters to your optimizer if you want them to be updated with gradient descent.

Upvotes: 3

Related Questions