Training on GPU

[gombru]

I saw that the code has this commented line:
# torch.cuda.set_device(args.gpu)
And by default it runs on CPU. If I uncomment the line and set
torch.cuda.set_device(0)
it still runs on CPU. And if I call
model.cuda()
after initializing the model it crashes saying it doesn't get the expected tensor types.
Is the code prepared to run in GPU? How can I do it?