Outside of an autograd function, the following code shows that b
requires_grad, as I’d expect to happen:
import numpy as np
import torch
from torch import nn, optim
N = 5
K = 3
a = nn.Parameter(torch.rand(N, K))
idxes = torch.from_numpy(np.random.choice(N, K, replace=True))
print('a.requires_grad', a.requires_grad)
b = a[idxes]
print('b.requires_grad', b.requires_grad)
print('type(a)', type(a))
print('type(b)', type(b))
output:
a.requires_grad True
b.requires_grad True
type(a) <class 'torch.nn.parameter.Parameter'>
type(b) <class 'torch.Tensor'>
however, if I run this inside an autograd function, b no longer requires_grad:
import numpy as np
import torch
from torch import nn, optim, autograd
def run2():
class my_function(autograd.Function):
@staticmethod
def forward(ctx, x):
print('x.requires_grad', x.requires_grad)
N, K = x.size()
idxes = torch.from_numpy(np.random.choice(N, K, replace=True))
z = x[idxes]
print('z.requires_grad', z.requires_grad)
return z
my_function = my_function.apply
N = 5
K = 3
# a = torch.rand(N, K, requires_grad=True)
a = torch.rand(N, K)
a = nn.Parameter(a)
my_function(a)
if __name__ == '__main__':
run2()
Output:
x.requires_grad True
z.requires_grad False
Why is the z in this case marked as not requiring grad? What can I do so that it does require grad, and so that gradients flowing through z will flow into x?