How to tell PyTorch which CUDA version to take?

Question

I have two version of CUDA installed on my Ubuntu 16.04 machine: 9.0 and 10.1. They are located in /usr/local/cuda-9.0 and /usr/local/10.1 respectively