Python Shell that has PyTorch Cuda Support

So, you just got nixos, and you want to pip install torch with cuda support. You create your venv, do your pip install, try everything, and you never get torch.cuda.is_available() to return true.

Turns out that all you need to do is source the correct libraries:

let
pkgs = import {};
in
pkgs.mkShellNoCC {
env = {
LD_LIBRARY_PATH = pkgs.lib.makeLibraryPath [
pkgs.stdenv.cc.cc
pkgs.libz
“/run/opengl-driver”
];
};
}

save that into a shell.nix, and run $ nix-shell in the same directory where you saved that file

hope this helps someone out there

technically you don’t even need libz to get that line of python to return true, but you will need it to do just about anything else