If you had an problems with th CUDA version :
You should follow these steps:
rm -r miniconda3
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
chmod +x Miniconda3-latest-Linux-x86_64.sh
source ~/.bashrc
or simply
conda activate
ssh lipn-gpus-a40
pip3 install torch torchvision torchaudio
srun python3 gpu_test.py
IT'S VERY IMPORTANT TO FOLLOW THESE STEPS IN THE RIGHT ORDER