Support Cuda for PyTorch in NixOS

with-cuda.nix

# Run with `nix-shell cuda-shell.nix`
{ pkgs ? import <nixpkgs> {} }:
pkgs.mkShell {
   name = "cuda-env-shell";
   buildInputs = with pkgs; [
     git gitRepo gnupg autoconf curl
     procps gnumake util-linux m4 gperf unzip
     cudatoolkit linuxPackages.nvidia_x11
     libGLU libGL
     xorg.libXi xorg.libXmu freeglut
     xorg.libXext xorg.libX11 xorg.libXv xorg.libXrandr zlib 
     ncurses5 stdenv.cc binutils
   ];
   shellHook = ''
      export CUDA_PATH=${pkgs.cudatoolkit}
      export LD_LIBRARY_PATH=${pkgs.lib.makeLibraryPath [ pkgs.stdenv.cc.cc.lib pkgs.linuxPackages.nvidia_x11 pkgs.ncurses5 pkgs.zlib ]}
      export EXTRA_LDFLAGS="-L/lib -L${pkgs.lib.makeLibraryPath [ pkgs.linuxPackages.nvidia_x11 ]}"
      export EXTRA_CCFLAGS="-I/usr/include"
   '';          
}