blob: 14560b06f87ce15969051de59d8ed911b77aa1db (
plain) (
blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
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)
''
|
