1 files changed, 20 insertions, 0 deletions
diff --git a/pkgs/development/python-modules/torch/mk-runtime-check.nix b/pkgs/development/python-modules/torch/mk-runtime-check.nix
new file mode 100644
index 0000000000000..61180a19aaba5
--- /dev/null
+++ b/pkgs/development/python-modules/torch/mk-runtime-check.nix
@@ -0,0 +1,20 @@
+{
+  cudaPackages,
+  feature,
+  libraries,
+  versionAttr,
+  pythonPackages,
+}:
+
+(cudaPackages.writeGpuTestPython.override { python3Packages = pythonPackages; })
+  {
+    inherit feature;
+    inherit libraries;
+    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)
+  ''