python3Packages.torchvision: added cudaSupport option (#132917)

Co-authored-by: Sandro <sandro.jaeckel@gmail.com>
2021-08-08 19:42:58 +02:00
parent 0d078fcdb2
commit 717538e908
2 changed files with 26 additions and 3 deletions
@@ -301,6 +301,11 @@ in buildPythonPackage rec {
  # Builds in 2+h with 2 cores, and ~15m with a big-parallel builder.
  requiredSystemFeatures = [ "big-parallel" ];

+  passthru = {
+    inherit cudaSupport;
+    cudaArchList = final_cudaArchList;
+  };
+
  meta = with lib; {
    description = "Open source, prototype-to-production deep learning platform";
    homepage    = "https://pytorch.org/";