Pytorch installed via pip does not pick up CUDA

Help
1

I have a simple shell.nix that picks up python3.9 from nix store and immediately enters Pythons virtualenv where I installed pytorch with CUDA support from pip (using Pytorch’s official instructions).

with import <nixpkgs> {};

mkShell {
  buildInputs = [
    stdenv.cc.cc.lib
    python39
    python39Packages.pip
  ];
  shellHook = ''
    # Add the library path for libstdc++.so.6
    export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${stdenv.cc.cc.lib}/lib"
    source .env/bin/activate
  '';
}

However it does not pick up CUDA:

$ python -c "import torch; print(torch.cuda.is_available())"
False

Any idea on how to troubleshoot this? I’m on CentOS 7 with nix installed via nix-portable. Normally my workflow involves conda environments, installing pytorch from pip with conda’s python works well.

1 Like
2

Did you ever get this to work? I’m running into exactly the same issue.

3