pyrox.dev
1{
2 cudaPackages,
3 feature,
4 libraries,
5 versionAttr,
6}:
7
8cudaPackages.writeGpuTestPython
9 {
10 inherit feature;
11 inherit libraries;
12 name = "${feature}Available";
13 }
14 ''
15 import torch
16 message = f"{torch.cuda.is_available()=} and {torch.version.${versionAttr}=}"
17 assert torch.cuda.is_available() and torch.version.${versionAttr}, message
18 print(message)
19 ''