cuda/pythonflake/flake.nix


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

{
  description = "Torch cuda flake using nix-community cachix";

  nixConfig = {
    extra-substituters = [
      "https://nix-community.cachix.org"
      "https://cuda-maintainers.cachix.org"
    ];
    extra-trusted-public-keys = [
      "nix-community.cachix.org-1:mB9FSh9qf2dCimDSUo8Zy7bkq5CX+/rkCWyvRCYg3Fs="
      "cuda-maintainers.cachix.org-1:0dq3bujKpuEPMCX6U4WylrUDZ9JyUG0VpVZa7CNfq5E="
    ];
  };

  outputs = {
    self,
    nixpkgs,
  }: let
    pkgs = import nixpkgs {
      system = "x86_64-linux";
      config.allowUnfree = true;
      config.cudaSupport = true;
    };
  in {
    devShell.x86_64-linux = with pkgs;
      mkShell {
        venvDir = "./.venv";
        buildInputs = [
          (pkgs.python3.withPackages (
            ps:
              with ps; [
                torch
              ]
          ))
          pkgs.virtualenv
          pkgs.python3Packages.venvShellHook
        ];
        postVenvCreation = ''
          unset SOURCE_DATE_EPOCH
        '';
      };
  };
}