1 files changed, 43 insertions, 0 deletions
diff --git a/cuda/pythonflake/flake.nix b/cuda/pythonflake/flake.nix
new file mode 100644
index 0000000..27cac81
--- /dev/null
+++ b/cuda/pythonflake/flake.nix
@@ -0,0 +1,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
+        '';
+      };
+  };
+}