Negative impact of Python versions on using Pytorch with Cuda11

Let us create conda virtual env with 3.6.4 or 3.8.5

  1. conda create -n myenv python=3.6.4 or conda create -n myenv python=3.8.5

  2. conda install pytorch cudatoolkit=11.0 -c pytorch

import torch
print(torch.cuda.get_device_name('cuda:0')) => 'GeForce RTX 3090'
size=100
X_cpu = torch.randn((size, size))
X_gpu = X_cpu.to(torch.device('cuda:0')) # (1)
torch.mm(X_gpu,X_gpu) # (2)

Bad news => The execution stalls at (1) or (2) depending on whether you are using python=3.6.4, python=3.8.5.

Good news => The execution does not stall if you are using python=3.7.1

Any idea why does this problem occur ?

cc @ngimel

1 possible answer(s) on “Negative impact of Python versions on using Pytorch with Cuda11

  1. Executing these two commands
    conda create -n myenv python=3.8.5 and conda install pytorch cudatoolkit=11.0 -c pytorch results in installing pytorch=1.5.1., (not pytorch=1.7.1.) as shown below.

    That’s not the case for me and also it seems that 1.7.1 is installed on your system:

    (myenv) datascienceadmin@limbo:~$ conda install pytorch cudatoolkit=11.0 -c pytorch
    [...]
    The following NEW packages will be INSTALLED:
    [...]
      pytorch            pytorch/linux-64::pytorch-1.7.1-py3.8_cuda11.0.221_cudnn8.0.5_0
    

    However, since 1.5.1 is found in your REPL, I guess you might have multiple PyTorch versions installed, where the 1.5.1 version is picked. This would also explain the reported “hang”, since 1.5.1 doesn’t ship with CUDA11.0 and the CUDA JIT would kick in, which might be slow.

    I would recommend to uninstall all conda, pip, and source builds in the current environment, make sure that import torch fails, and reinstall 1.7.1.