python3Packages.torch-bin: gpuChecks -> tests.tester-<name>.gpuCheck

pkgs/development/python-modules/torch/bin.nix

@@ -121,7 +121,10 @@ buildPythonPackage {
 
   pythonImportsCheck = [ "torch" ];
 
-  passthru.gpuChecks.cudaAvailable = callPackage ./test-cuda.nix { torch = torch-bin; };
+  passthru.tests = callPackage ./tests.nix {
+    torchWithCuda = torch-bin;
+    torchWithRocm = torch-bin;
+  };
 
   meta = {
     description = "PyTorch: Tensors and Dynamic neural networks in Python with strong GPU acceleration";

pkgs/development/python-modules/torch/mk-runtime-check.nix

@@ -0,0 +1,19 @@
+{
+  cudaPackages,
+  feature,
+  torch,
+  versionAttr,
+}:
+
+cudaPackages.writeGpuTestPython
+  {
+    inherit feature;
+    libraries = [ torch ];
+    name = "${feature}Available";
+  }
+  ''
+    import torch
+    message = f"{torch.cuda.is_available()=} and {torch.version.${versionAttr}=}"
+    assert torch.cuda.is_available() and torch.version.${versionAttr}, message
+    print(message)
+  ''

pkgs/development/python-modules/torch/tests.nix

@@ -1,3 +1,21 @@
-{ callPackage }:
+{
+  callPackage,
+  torchWithCuda,
+  torchWithRocm,
+}:
 
-callPackage ./gpu-checks.nix { }
+{
+  # To perform the runtime check use either
+  # `nix run .#python3Packages.torch.tests.tester-cudaAvailable` (outside the sandbox), or
+  # `nix build .#python3Packages.torch.tests.tester-cudaAvailable.gpuCheck` (in a relaxed sandbox)
+  tester-cudaAvailable = callPackage ./mk-runtime-check.nix {
+    feature = "cuda";
+    versionAttr = "cuda";
+    torch = torchWithCuda;
+  };
+  tester-rocmAvailable = callPackage ./mk-runtime-check.nix {
+    feature = "rocm";
+    versionAttr = "hip";
+    torch = torchWithRocm;
+  };
+}