The following autograd function attempts to calculate some auxiliary loss and backprop it. However I get an error element 0 of tensors does not require grad and does not have a grad_fn, even though x requires a gradient. Why?
import numpy as np

import torch
from torch import nn, optim, autograd


def run3():
    class my_function(autograd.Function):
        @staticmethod
        def forward(ctx, x):
            ctx.save_for_backward(x)
            return x

        @staticmethod
        def backward(ctx, grad_output):
            (x,) = ctx.saved_variables
            print('x.requires_grad', x.requires_grad)
            aux_loss = x.sum()
            aux_loss.backward(retain_graph=True)
            x.backward(grad_output)


    my_function = my_function.apply
    N = 5
    K = 3

    a = torch.rand(N, K)
    a = nn.Parameter(a)
    out = my_function(a)
    loss = out.sum()
    print('loss', loss)
    loss.backward()


if __name__ == '__main__':
    run3()

Output:
x.requires_grad True
...
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

Edit: it turns out that aux_loss doesnt have require_grad set. So, I reckon this question relates to A tensor formed by indexing a tensor inside an autograd function does not require grad 3 . Seems something works differently inside of autograd functions to normal somehow?