mirror of
https://github.com/NixOS/nixpkgs.git
synced 2025-01-01 02:23:54 +00:00
21 lines
453 B
Nix
21 lines
453 B
Nix
{
|
|
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)
|
|
''
|