diff options
Diffstat (limited to 'pkgs/development/python-modules/torch/mk-runtime-check.nix')
-rw-r--r-- | pkgs/development/python-modules/torch/mk-runtime-check.nix | 20 |
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) + '' |