Hello I am writing a small pytorch example with a simple NN. The program runs fine if I declare
dtype = torch.FloatTensor
#dtype = torch.cuda.FloatTensor # Uncomment this to run on GPU
The code currently runs great with the CPU option. However, as soon as I uncomment and switch to the GPU option, the code hangs when I call the function
x = Variable(x.type(dtype), requires_grad=True)
The code essentially gets stuck there and stop working.
I have tried to debug this problem. I have confirmed that I have cuda installed
~ $ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2015 NVIDIA Corporation
Built on Tue_Aug_11_14:27:32_CDT_2015
Cuda compilation tools, release 7.5, V7.5.17
I have confirmed that pytorch can see the cuda
torch.cuda.is_available()
True
I assume I am just missing something small, any help would be appreciated.
The full code can be seen here :