Can't install GPU-enabled Pytorch in Conda environment from environment.yml

user_1730417 0
Viewed 4752

I'm on Ubuntu 20.04 LTS with CUDA 11.1 installed (and working, with PATH and LD_LIBRARY_PATH configured correctly), and I'm trying to define a reusable conda environment (i.e., in an environment.yml file) that successfully installs PyTorch with CUDA support.

However, when I use the environment file, I get a message that Torch wasn't compiled with CUDA support:

Python 3.8.10 | packaged by conda-forge | (default, May 11 2021, 07:01:05)
[GCC 9.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> device = torch.device("cuda:0")
>>> t = torch.tensor(device=device, data=[0,1,2,3])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/jdr2160/anaconda3/envs/foo/lib/python3.8/site-packages/torch/cuda/__init__.py", line 166, in _lazy_init
    raise AssertionError("Torch not compiled with CUDA enabled")
AssertionError: Torch not compiled with CUDA enabled

My environment.yml is pretty bare-bones:

name: foo
channels:
  - conda-forge
  - nvidia
  - pytorch
dependencies:
  - cudatoolkit=11.1
  - python=3.8
  - pytorch

When I create an 'empty' python 3.8 environment and install the Conda packages from the command line instead of from an environment file, everything works fine:

$ conda env create --name bar python=3.8
...
$ conda activate bar
$ conda install pytorch cudatoolkit=11.1 -c pytorch -c nvidia
...
$ python
Python 3.8.10 | packaged by conda-forge | (default, May 11 2021, 07:01:05)
[GCC 9.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> device = torch.device("cuda:0")
>>> t = torch.tensor(device=device, data=[0,1,2,3])
>>>

Can anyone tell what's going on here? It seems that Conda doesn't see the cudatoolkit=11.1 dependency while installing PyTorch from the environment file, but I have no idea how to fix it.

pytorch conda python
2 Answers

Just a few minutes after posting this question, I was able to figure out the solution. It turns out that it has to do with prioritizing Conda channels. The solution (which isn't well-documented by Anaconda) is to specify the correct channel for cudatoolkit and pytorch in environment.yml:

name: foo
channels:
  - conda-forge
  - nvidia
  - pytorch
dependencies:
  - nvidia::cudatoolkit=11.1
  - python=3.8
  - pytorch::pytorch
user_17304170

For conda version: 4.10.3, 4.11.0

  • conda --version to get your version
  • conda update -n base -c defaults conda to update your conda

This one worked for me:

in enviroment.yaml

name: nlp

channels:
  - pytorch

dependencies:
  - python=3.9
  - numpy=1.21.5
  - pandas=1.3.5
  - spacy=3.2.1
  - tensorflow=2.6.0
  - pytorch=1.10.1
  - cudatoolkit=11.3

in terminal

conda env create --file environment.yaml
conda activate nlp # use your env name from enviroment.yaml
python main.py

in main.py

import numpy as np
import pandas as pd
import spacy
import tensorflow as tf
import torch

print(f'np: {np.__version__}')
print(f'pd: {pd.__version__}')
print(f'spacy: {spacy.__version__}')
print(f'tf: {tf.__version__}')
print(f'torch: {torch.__version__}')

print(f'cuda enable: {torch.cuda.is_available()}')
print(f'current_device: {torch.cuda.current_device()}')
print(f'device: {torch.cuda.device(0)}')
print(f'device_count: {torch.cuda.device_count()}')
print(f'get_device_name: {torch.cuda.get_device_name(0)}')

output

np: 1.21.5
pd: 1.3.5
spacy: 3.2.1
tf: 2.6.0
torch: 1.10.1
cuda enable: True
current_device: 0
device: <torch.cuda.device object at 0x0000015156785EB0>
device_count: 1
get_device_name: NVIDIA GeForce GTX 1650 Ti
user_139479310
Related