A tensor formed by indexing a tensor inside an autograd function does not require grad

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?

Similar to your other question, you need with torch.enable_grad(): if you want to record the graph within the forward or backward of a autograd.Function.

Best regards

Thomas

1 Like