Hi @albanD,
I am using the following implementation, but the gradient is  None w.r.t inputs.
import torch
import torch.nn as nn
from torch.autograd import Variable

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        pass

    def forward(self, x):
        self.input = []
        x = x.clone().detach().requires_grad_(True)
        for i in range(2):          
            x = x * 2
            # detach from previous history
            x = x.clone().detach().requires_grad_(True)
            self.input.append(x)
        return x

    def backward(self, p):
        
        for input in self.input:
            print(input)
            print(torch.autograd.grad(loss, input, allow_unused=True))

model = Net()
criterion = torch.nn.MSELoss()

"""
x = 2 # dl/ax == -4
a = 2 * x # dl/da == dl/db * db/da == -2
b = 2 * a # dl/db == -1
l = 10 - b # dl/dl == 1
"""

inp = torch.tensor(2.0)
output = model(inp) # predcicted output

output = output.clone().detach().requires_grad_(True)
loss = (10 - output).sum()
model.backward(loss)


Output should be like this right?

"""
x = 2 # dl/ax == -4
a = 2 * x # dl/da == dl/db * db/da == -2
b = 2 * a # dl/db == -1
l = 10 - b # dl/dl == 1
"""

need help. Thanks in advance.