1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
|
import torch
@torch.jit.script
def fn(x, scale, shift):
return scale * x / shift
@torch.jit.script
def recurrent(x, scale, shift):
y = x
for i in range(100):
y = fn(y, scale, shift)
return y
x = torch.randn(2, 2, device="cuda")
scale = torch.randn(2, 2, device="cuda", requires_grad=True)
shift = torch.randn(2, 2, device="cuda", requires_grad=True)
inputs = [x, scale, shift]
out = recurrent(x, scale, shift)
recurrent.graph_for(x, scale, shift)
import torch
@torch.jit.script
def recurrent_scaleshift(x, scale, shift):
y = x
for i in range(64):
y = scale * y + shift
return y
x = torch.randn(2, 2, device="cuda")
scale = torch.randn(2, 2, device="cuda", requires_grad=True)
shift = torch.randn(2, 2, device="cuda", requires_grad=True)
inputs = [x, scale, shift]
out = recurrent_scaleshift(x, scale, shift)
recurrent_scaleshift.graph_for(x, scale, shift)
import torch
x = torch.tensor([])
x.requires_grad = True
x.mean().backward() # no error triggered
x = x.cuda()
x.mean().backward()
|
