pyrox.dev
1{
2 cudaPackages,
3 feature ? null,
4 lib,
5 libraries,
6 name ? if feature == null then "torch-compile-cpu" else "torch-compile-${feature}",
7 stdenv,
8}:
9let
10 deviceStr = if feature == null then "" else '', device="cuda"'';
11in
12cudaPackages.writeGpuTestPython
13 {
14 inherit name feature libraries;
15 makeWrapperArgs = [
16 "--suffix"
17 "PATH"
18 ":"
19 "${lib.getBin stdenv.cc}/bin"
20 ];
21 }
22 ''
23 import torch
24
25
26 @torch.compile
27 def opt_foo2(x, y):
28 a = torch.sin(x)
29 b = torch.cos(y)
30 return a + b
31
32
33 print(
34 opt_foo2(
35 torch.randn(10, 10${deviceStr}),
36 torch.randn(10, 10${deviceStr})))
37 ''