1 files changed, 6 insertions, 1 deletions
diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix
index d5d7e823bed7c..9597a047bdb48 100644
--- a/pkgs/development/python-modules/torch/default.nix
+++ b/pkgs/development/python-modules/torch/default.nix
@@ -24,6 +24,10 @@
   mpi,
   buildDocs ? false,
 
+  # tests.cudaAvailable:
+  callPackage,
+  torchWithCuda,
+
   # Native build inputs
   cmake,
   symlinkJoin,
@@ -639,11 +643,12 @@ buildPythonPackage rec {
       rocmSupport
       rocmPackages
       ;
+    cudaCapabilities = if cudaSupport then supportedCudaCapabilities else [ ];
     # At least for 1.10.2 `torch.fft` is unavailable unless BLAS provider is MKL. This attribute allows for easy detection of its availability.
     blasProvider = blas.provider;
     # To help debug when a package is broken due to CUDA support
     inherit brokenConditions;
-    cudaCapabilities = if cudaSupport then supportedCudaCapabilities else [ ];
+    tests = callPackage ./tests.nix { };
   };
 
   meta = {