From 19bfd98a971770c03eb957a550cfff27e41569c9 Mon Sep 17 00:00:00 2001
From: DavHau <hsngrmpf+github@gmail.com>
Date: Wed, 13 Sep 2023 13:50:19 +0200
Subject: [PATCH] fix(overlays/torch): use autoAddOpenGLRunpathHook

---
 overrides/python/torch/default.nix | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/overrides/python/torch/default.nix b/overrides/python/torch/default.nix
index e1247637..97676fd3 100644
--- a/overrides/python/torch/default.nix
+++ b/overrides/python/torch/default.nix
@@ -6,12 +6,15 @@
   # stripping doesn't reduce the file size much, and it takes a long time
   mkDerivation.dontStrip = true;
 
+  # use the autoAddOpenGLRunpathHook to add /run/opengl-driver/lib to the RPATH
+  #   of all ELF files
+  deps = {nixpkgs, ...}: {
+    inherit (nixpkgs.cudaPackages) autoAddOpenGLRunpathHook;
+  };
+  mkDerivation.nativeBuildInputs = [
+    config.deps.autoAddOpenGLRunpathHook
+  ];
+
   # this file is patched manually, so ignore it in autoPatchelf
   env.autoPatchelfIgnoreMissingDeps = ["libcuda.so.1"];
-  # patch the rpath so libcuda.so.1 can be found at /run/opengl-driver/lib
-  env.cudaPatchPhase = ''
-    patchelf $out/${config.deps.python.sitePackages}/torch/lib/libcaffe2_nvrtc.so \
-      --add-rpath /run/opengl-driver/lib
-  '';
-  mkDerivation.postPhases = ["cudaPatchPhase"];
 }