Pytorch installed via pip does not pick up CUDA

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 = [
  shellHook = ''
    # Add the library path for
    export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${}/lib"
    source .env/bin/activate

However it does not pick up CUDA:

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

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

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