{
  description = "Pytorch and stuff";

  # Specifies other flakes that this flake depends on.
  inputs = {
    devshell.url = "github:numtide/devshell";
    utils.url = "github:numtide/flake-utils";
    nixpkgs.url = "github:nixos/nixpkgs/nixos-21.11";
  };

  # Function that produces an attribute set.
  # Its function arguments are the flakes specified in inputs.
  # The self argument denotes this flake.
  outputs = inputs@{ self, nixpkgs, utils, ... }:
    (utils.lib.eachSystem [ "x86_64-linux" ] (system:
      let
        pkgs = (import nixpkgs {
          inherit system;
          config = {
            # For CUDA.
            allowUnfree = true;
            # Enables CUDA support in packages that support it.
            cudaSupport = true;
          };
        });
      in rec {
        # Executed by `nix build .#<name>`
        # Ignore this, it was just for testing.
        packages = utils.lib.flattenTree {
          hello = pkgs.hello;
        };

        # Executed by `nix build .`
        defaultPackage = packages.hello;
        # defaultPackage = pkgs.callPackage ./default.nix { };

        # Executed by `nix develop`
        devShell = with pkgs; mkShell {
          buildInputs = [
            python39 # numba-0.54.1 not supported for interpreter python3.10
          ] ++ (with python39.pkgs; [
            inflect
            librosa
            pip
            pytorch-bin
            unidecode
          ]) ++ (with cudaPackages; [
            cudatoolkit
          ]);

          shellHook = ''
            export CUDA_PATH=${pkgs.cudatoolkit}
          '';
        };
      }
    ));
}

In reply to @yuu:matrix.org
anyone who has a flake for pytorch + cuda development? or it could be a shell.nix too

In reply to @yuu:matrix.org

{
  description = "Pytorch and stuff";

  # Specifies other flakes that this flake depends on.
  inputs = {
    devshell.url = "github:numtide/devshell";
    utils.url = "github:numtide/flake-utils";
    nixpkgs.url = "github:nixos/nixpkgs/nixos-21.11";
  };

  # Function that produces an attribute set.
  # Its function arguments are the flakes specified in inputs.
  # The self argument denotes this flake.
  outputs = inputs@{ self, nixpkgs, utils, ... }:
    (utils.lib.eachSystem [ "x86_64-linux" ] (system:
      let
        pkgs = (import nixpkgs {
          inherit system;
          config = {
            # For CUDA.
            allowUnfree = true;
            # Enables CUDA support in packages that support it.
            cudaSupport = true;
          };
        });
      in rec {
        # Executed by `nix build .#<name>`
        # Ignore this, it was just for testing.
        packages = utils.lib.flattenTree {
          hello = pkgs.hello;
        };

        # Executed by `nix build .`
        defaultPackage = packages.hello;
        # defaultPackage = pkgs.callPackage ./default.nix { };

        # Executed by `nix develop`
        devShell = with pkgs; mkShell {
          buildInputs = [
            python39 # numba-0.54.1 not supported for interpreter python3.10
          ] ++ (with python39.pkgs; [
            inflect
            librosa
            pip
            pytorch-bin
            unidecode
          ]) ++ (with cudaPackages; [
            cudatoolkit
          ]);

          shellHook = ''
            export CUDA_PATH=${pkgs.cudatoolkit}
          '';
        };
      }
    ));
}

Someone S: hi there ^-^ i just wanted a minimal flake.nix with pytorch and cuda enabled. i did it with that flake

{
  description = "Pytorch and stuff";

  # Specifies other flakes that this flake depends on.
  inputs = {
    devshell.url = "github:numtide/devshell";
    utils.url = "github:numtide/flake-utils";
    nixpkgs.url = "github:nixos/nixpkgs/nixos-21.11";
  };

  # Function that produces an attribute set.
  # Its function arguments are the flakes specified in inputs.
  # The self argument denotes this flake.
  outputs = inputs@{ self, nixpkgs, utils, ... }:
    (utils.lib.eachSystem [ "x86_64-linux" ] (system:
      let
        pkgs = (import nixpkgs {
          inherit system;
          config = {
            # For CUDA.
            allowUnfree = true;
            # Enables CUDA support in packages that support it.
            cudaSupport = true;
          };
        });
      in rec {
        # Executed by `nix build .#<name>`
        # Ignore this, it was just for testing.
        packages = utils.lib.flattenTree {
          hello = pkgs.hello;
        };

        # Executed by `nix build .`
        defaultPackage = packages.hello;
        # defaultPackage = pkgs.callPackage ./default.nix { };

        # Executed by `nix develop`
        devShell = with pkgs; mkShell {
          buildInputs = [
            python39 # numba-0.54.1 not supported for interpreter python3.10
          ] ++ (with python39.pkgs; [
            inflect
            librosa
            pip
            pytorch-bin
            unidecode
          ]) ++ (with cudaPackages; [
            cudatoolkit
          ]);

          shellHook = ''
            export CUDA_PATH=${pkgs.cudatoolkit}
          '';
        };
      }
    ));
}

and direnv + nix-direnv .envrc

use flake

👍️

I'll try to send some of the shells that I use a little later.
But meanwhile a few links you might find useful, if you haven't seen them yet:

• zimbatm's nixpkgs-unfree: https://discourse.nixos.org/t/announcing-nixpkgs-unfree/17505
  The idea is to use inputs.nixpkgs.url = github:numtide/nixpkgs-unfree/$branch as input, which exposes nixpkgs with the unfree already enabled and tracks NixOS/nixpkgs automatically. I'm using smth like this for the flake registry in my configuration.nix, so that nixpkgs#python3Packages.pytorch resolves into the pre-cached unfree and cuda-enabled pytorch for me, instead of throwing an error about NIXPKGS_ALLOW_UNFREE=1
• It's somewhat experimental, but there's now cachix with prebuilt cuda-enabled sci-comp packages, you can find more at https://nixos.wiki/wiki/CUDA. The cache is populated automatically for the last release, unstable, and master, much like nixpkgs-unfree does it. There's a slight delay between a branch update and the cache, but probably not more than a day
• I prefer using source-based pytorch instead of the wheel-based pytorch-bin. The way I use cuda-enabled packages is essentially just importing nixpkgs with { config = { allowUnfree = true; cudaSupport = true; }; } and using directly pytorch, jax, blender, etc. And I'd rather consume nixpkgs-unfree as a flake input, than import it myself in a project.