Pytorch and CUDA: "Torch not compiled with CUDA enabled"

Help
1

Hello, I’m new on the NixOS world, but I think it can be very useful for setting up data-science environments.

I wanted to create a default.nix file that could provide all devs the same environment, but I get always the error: “Torch not compiled with CUDA enabled”

After reading the tutorials and the docs, I came up with this:

{ sources ? import ./nix/sources.nix
, pkgs ? import sources.nixpkgs {}
}:

pkgs.mkShell {
  buildInputs = [
    pkgs.python38
    pkgs.cudaPackages.cudatoolkit_10
    pkgs.python38Packages.pytorch
    pkgs.python38Packages.unidecode
    pkgs.python38Packages.inflect
    pkgs.python38Packages.librosa
    pkgs.python38Packages.pip
  ];

  shellHook = ''
    echo "You are now using a NIX environment"
    export CUDA_PATH=${pkgs.cudatoolkit}
  '';
}

I’m using branch 20.9

"nixpkgs": {
        "branch": "nixos-20.09",
        "description": "Nix Packages collection",
        "homepage": "",
        "owner": "NixOS",
        "repo": "nixpkgs",
        "rev": "88f00e7e12d2669583fffd3f33aae01101464386",
        "sha256": "0972lcah2wm1j7ab5acnpn1il68q90cdqhvq1vj4nlnygnwzhcfr",
        "type": "tarball",
        "url": "https://github.com/NixOS/nixpkgs/archive/88f00e7e12d2669583fffd3f33aae01101464386.tar.gz",
        "url_template": "https://github.com/<owner>/<repo>/archive/<rev>.tar.gz"
    },

But when I run:

nix-shell
python
>>> import torch
>>> torch.cuda.current_device()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/nix/store/9l5ganwsbn8cir7skq1af8y1jf4przmx-python3.8-pytorch-1.6.0/lib/python3.8/site-packages/torch/cuda/__init__.py", line 384, in current_device
    _lazy_init()
  File "/nix/store/9l5ganwsbn8cir7skq1af8y1jf4przmx-python3.8-pytorch-1.6.0/lib/python3.8/site-packages/torch/cuda/__init__.py", line 186, in _lazy_init
    _check_driver()
  File "/nix/store/9l5ganwsbn8cir7skq1af8y1jf4przmx-python3.8-pytorch-1.6.0/lib/python3.8/site-packages/torch/cuda/__init__.py", line 61, in _check_driver
    raise AssertionError("Torch not compiled with CUDA enabled")
AssertionError: Torch not compiled with CUDA enabled

How can I make this work?

Thanks!

1 Like
2

You need to enable CUDA when importing nixpkgs, since the default is to build PyTorch without CUDA support (since CUDA is non-free). E.g.:

import sources.nixpkgs {
  config = {
    allowUnfree = true;
    cudaSupport = true;
  };
}

cudaSupport enables CUDA for all packages that support this option.

Since unfree packages are not built by Hydra, CUDA-enabled PyTorch is not in the binary cache. If you want to avoid the long-ish build (depending on your hardware), you can also use pytorch-bin (in your case pkgs.python38Packages.pytorch-bin). pytorch-bin uses the upstream PyTorch.org builds, patched to work with libraries in the Nix store.

4 Likes
3

thanks a lot @danieldk ! Very useful information, I’m going to try it right away.

1 Like