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