pyrox.dev
1{ callPackage }:
2
3rec {
4 # To perform the runtime check use either
5 # `nix run .#python3Packages.torch.tests.tester-cudaAvailable` (outside the sandbox), or
6 # `nix build .#python3Packages.torch.tests.tester-cudaAvailable.gpuCheck` (in a relaxed sandbox)
7 tester-cudaAvailable = callPackage ./mk-runtime-check.nix {
8 feature = "cuda";
9 versionAttr = "cuda";
10 libraries = ps: [ ps.torchWithCuda ];
11 };
12 tester-rocmAvailable = callPackage ./mk-runtime-check.nix {
13 feature = "rocm";
14 versionAttr = "hip";
15 libraries = ps: [ ps.torchWithRocm ];
16 };
17
18 compileCpu = tester-compileCpu.gpuCheck;
19 tester-compileCpu = callPackage ./mk-torch-compile-check.nix {
20 feature = null;
21 libraries = ps: [ ps.torch ];
22 };
23 tester-compileCuda = callPackage ./mk-torch-compile-check.nix {
24 feature = "cuda";
25 libraries = ps: [ ps.torchWithCuda ];
26 };
27 tester-compileRocm = callPackage ./mk-torch-compile-check.nix {
28 feature = "rocm";
29 libraries = ps: [ ps.torchWithRocm ];
30 };
31
32 mnist-example = callPackage ./mnist-example { };
33}