Add pytorch-lightning from unstable while keeping torch-bin from stable.

Author: matthid

nix/dependencies.nix

@@ -1,4 +1,7 @@
 { pkgs ? import <nixpkgs> {
+    overlays = [
+      (import ./overlay/replace-torch.nix { })
+    ];
     config = {
       allowUnfree = true;
       cudaSupport = true;
@@ -7,9 +10,21 @@
 , lib ? pkgs.lib
 , my_python ? pkgs.python3
 , cudatoolkit ? pkgs.cudaPackages.cudatoolkit
-, }:
+, unstable_pkgs ? import (fetchTarball "https://github.com/NixOS/nixpkgs/archive/nixos-unstable.tar.gz") {
+    overlays = [
+      (import ./overlay/replace-torch.nix {
+          do_replace = true;
+          replacement_torch = my_python.pkgs.torch;
+          replacement_torchvision = my_python.pkgs.torchvision;
+          replacement_torchaudio = my_python.pkgs.torchaudio;
+          replacement_python = my_python;
+      })
+    ];
+  }
+}:
 let
   python_packages = my_python.pkgs;
+  unstable_python_packages = unstable_pkgs.my_python.pkgs;
 in {
   pkgs = pkgs;
   lib = lib;
@@ -18,6 +33,7 @@ in {
   dependencies = with pkgs; [
     python_packages.numpy
     python_packages.torch-bin
+    unstable_python_packages.pytorch-lightning
     cudatoolkit
   ];
 }

nix/overlay/replace-torch.nix

@@ -0,0 +1,28 @@
+{ do_replace ? false
+, replacement_torch ? false
+, replacement_torchvision ? false
+, replacement_torchaudio ? false
+, replacement_python ? false }:
+final: prev:
+let
+  real_python = if do_replace then replacement_python else prev.python311;
+  real_torch = if do_replace then replacement_torch else real_python.pkgs.torch-bin;
+  real_torchvision = if do_replace then replacement_torchvision else real_python.pkgs.torchvision-bin;
+  real_torchaudio = if do_replace then replacement_torchaudio else real_python.pkgs.torchaudio-bin;
+  real_python311 = real_python.override {
+    packageOverrides = final_: prev_: {
+      torch = real_torch;
+      torchvision = real_torchvision;
+      torchaudio = real_torchaudio;
+      pytorch-lightning = prev_.pytorch-lightning.override {
+        torch = real_torch;
+      };
+      tensorboardx = prev_.tensorboardx.override {
+        torch = real_torch;
+      };
+    };
+  };
+in {
+  python311 = real_python311;
+  my_python = real_python311;
+}