Yup. You have to first do the CUDA binaries installation.
For UBUNTU:
https://developer.nvidia.com/cuda-downloads?target_os=Linux&target_arch=x86_64&target_distro=Ubuntu&target_version=1604&target_type=debnetwork 381
After that go to:
http://pytorch.org 134
and install required packages. it will work like charm.
Also, you can do the simple test as described by Soumith.
import torch
torch.cuda.is_available()
if the above lines return True then you have Cuda working else not.