As recommended by you I have built PyTorch from source, I am able to train on Tesla k40c but when trying to run on GTX 1080 I get the following error:
GeForce GTX 1080 with CUDA capability sm_61 is not compatible with the current PyTorch installation.
The current PyTorch install supports CUDA capabilities sm_35.
If you want to use the GeForce GTX 1080 GPU with PyTorch, please check the instructions at https://pytorch.org/get-started/locally/
warnings.warn(incompatible_device_warn.format(device_name, capability, " ".join(arch_list), device_name))
Traceback (most recent call last):
File "<string>", line 1, in <module>
RuntimeError: CUDA error: no kernel image is available for execution on the device