Dear All,
I have installed Nvidia Cuda 9.0 toolkit with Cudnn to my ubuntu machine.
I have installed pytorch when i am trying to check for gpu usage by running the below code -
Code-
import torch
print(torch.rand(2,3).cuda())
I am getting the below error:

RuntimeError                              Traceback (most recent call last)
 in ()
----> 1 print(torch.rand(2,3).cuda())
~/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/_utils.py in cuda(self, device, async)
67         else:
68             new_type = getattr(torch.cuda, self.class.name)
—> 69             return new_type(self.size()).copy(self, async)
70
71
~/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/cuda/init.py in _lazy_new(cls, *args, **kwargs)
385     # We need this method only for lazy init, so we can remove it
386     del _CudaBase.new
–> 387     return super(_CudaBase, cls).new(cls, *args, **kwargs)
388
389
RuntimeError: cuda runtime error (2) : out of memory at /opt/conda/conda-bld/pytorch_1518244421288/work/torch/lib/THC/generic/THCStorage.cu:58
I think pytorch is not communicating with the Nvidia GPU, please advise.
Regards
Saurabh Jha