about summary refs log tree commit diff
path: root/pkgs/development/python-modules/pytorch/default.nix
diff options
context:
space:
mode:
Diffstat (limited to 'pkgs/development/python-modules/pytorch/default.nix')
-rw-r--r--pkgs/development/python-modules/pytorch/default.nix5
1 files changed, 5 insertions, 0 deletions
diff --git a/pkgs/development/python-modules/pytorch/default.nix b/pkgs/development/python-modules/pytorch/default.nix
index 35eb79d8b2d3..0de0015ab1ec 100644
--- a/pkgs/development/python-modules/pytorch/default.nix
+++ b/pkgs/development/python-modules/pytorch/default.nix
@@ -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/";