From d779b1adcb880da5e4eef661548095f9fa77b8c8 Mon Sep 17 00:00:00 2001
From: Thiago Franco de Moraes <totonixsame@gmail.com>
Date: Mon, 19 Apr 2021 21:48:55 -0300
Subject: [PATCH 1/2] Added an override to torchvision

---
 overrides.nix | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/overrides.nix b/overrides.nix
index 9d9a502..388b626 100644
--- a/overrides.nix
+++ b/overrides.nix
@@ -1459,6 +1459,28 @@ self: super:
     )
     { };
 
+  torchvision = lib.makeOverridable
+    ({ enableCuda ? false
+     , cudatoolkit ? pkgs.cudatoolkit_10_1
+     , pkg ? super.torchvision
+     }: pkg.overrideAttrs (old: {
+
+      # without that autoPatchelfHook will fail because cudatoolkit is not in LD_LIBRARY_PATH
+      autoPatchelfIgnoreMissingDeps = true;
+      buildInputs = (old.buildInputs or [ ])
+        ++ [ self.torch ]
+        ++ lib.optionals enableCuda [
+        cudatoolkit
+      ];
+      preConfigure =
+        if (enableCuda) then ''
+          export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${self.torch}/${self.python.sitePackages}/torch/lib:${lib.makeLibraryPath [ cudatoolkit "${cudatoolkit}" ]}"
+        '' else ''
+          export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${self.torch}/${self.python.sitePackages}/torch/lib"
+        '';
+    }))
+    { };
+
   typeguard = super.typeguard.overridePythonAttrs (old: {
     postPatch = ''
       substituteInPlace setup.py \

From 1b3d7f4f36554ea8e497beb66921d85057873631 Mon Sep 17 00:00:00 2001
From: Thiago Franco de Moraes <totonixsame@gmail.com>
Date: Wed, 21 Apr 2021 19:39:40 -0300
Subject: [PATCH 2/2] Added typing-extensions to torch propagatedBuildInputs

---
 overrides.nix | 1 +
 1 file changed, 1 insertion(+)

diff --git a/overrides.nix b/overrides.nix
index 388b626..293e9e3 100644
--- a/overrides.nix
+++ b/overrides.nix
@@ -1454,6 +1454,7 @@ self: super:
         propagatedBuildInputs = [
           self.numpy
           self.future
+          self.typing-extensions
         ];
       })
     )