torch-python-lib-rel-path.diff

diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix
index 4e3d54f4caa4..00d6b2727981 100644
--- a/pkgs/development/python-modules/torch/default.nix
+++ b/pkgs/development/python-modules/torch/default.nix
@@ -232,7 +232,12 @@ buildPythonPackage rec {
   };

   patches =
-    lib.optionals cudaSupport [ ./fix-cmake-cuda-toolkit.patch ]
+    [
+      # Allow setting PYTHON_LIB_REL_PATH with an environment variable.
+      # https://github.com/pytorch/pytorch/pull/128419
+      ./passthrough-python-lib-rel-path.patch
+    ]
+    ++ lib.optionals cudaSupport [ ./fix-cmake-cuda-toolkit.patch ]
     ++ lib.optionals (stdenv.isDarwin && stdenv.isx86_64) [
       # pthreadpool added support for Grand Central Dispatch in April
       # 2020. However, this relies on functionality (DISPATCH_APPLY_AUTO)
@@ -373,6 +378,10 @@ buildPythonPackage rec {
   USE_SYSTEM_NCCL = USE_NCCL;
   USE_STATIC_NCCL = USE_NCCL;

+  # Set the correct Python library path, broken since
+  # https://github.com/pytorch/pytorch/commit/3d617333e
+  PYTHON_LIB_REL_PATH = "${placeholder "out"}/${python.sitePackages}";
+
   # Suppress a weird warning in mkl-dnn, part of ideep in pytorch
   # (upstream seems to have fixed this in the wrong place?)
   # https://github.com/intel/mkl-dnn/commit/8134d346cdb7fe1695a2aa55771071d455fae0bc
diff --git a/pkgs/development/python-modules/torch/passthrough-python-lib-rel-path.patch b/pkgs/development/python-modules/torch/passthrough-python-lib-rel-path.patch
new file mode 100644
index 000000000000..629a0495c00e
--- /dev/null
+++ b/pkgs/development/python-modules/torch/passthrough-python-lib-rel-path.patch
@@ -0,0 +1,12 @@
+diff --git a/tools/setup_helpers/cmake.py b/tools/setup_helpers/cmake.py
+index 5481ce46031c..d50d9d547399 100644
+--- a/tools/setup_helpers/cmake.py
++++ b/tools/setup_helpers/cmake.py
+@@ -231,6 +231,7 @@ def generate(
+                     "SELECTED_OP_LIST",
+                     "TORCH_CUDA_ARCH_LIST",
+                     "TRACING_BASED",
++                    "PYTHON_LIB_REL_PATH",
+                 )
+             }
+         )