Hi @albanD,
Thank you for you replay.
I am trying to save the input to each layer and calculate the gradient w.r.t the input to each layer. Suppose we have a three layer in the network.
def model(x):
    a = w1 * x # layer 1
    b = w2 * a # layer 2
    c = w3 * b # layer 3
    return c

for epoch in epochs:
    # now i will train the model and update the weights w1, w2, w3. i.e
    for x in train_loader:
        optimizer.zero_grad()
        out = model(x)
        loss = 10 - y
        loss.backward()
        optimizers.step()

    # validate model
    ....
    ...

    # In this step I don't want to update the w1, w2, w3 I only want to 
    #calculate the gradient with respect to input to each layer 
    # (output of privious layer including) i.e dl/dx, dl/da, dl/db, dl/dc.

    # save this gradient and average it after all epochs is complete.
    # and gradient w.r.t inputs has be calculated by single input only (batch_size==1)
    for batch in train_loader:
        for image in batch:
            calculate_gradient for inputs to each layer; 
            save(gradient)


I hope you understand my problem. sorry for the poor explanation.currently I am doing this but I don’t know I am doing it right.
import torch
import torch.nn as nn
from torch.autograd import Variable

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.layers = nn.ModuleList([
            nn.Linear(4, 8),
            nn.Linear(8, 16),
            nn.Linear(16, 2),
        ])

    def forward(self, x):
        self.output = []
        self.input = []
        for layer in self.layers:
            # detach from previous history
            x = Variable(x.data, requires_grad=True)
            self.input.append(x)

            # compute output
            x = layer(x)
            self.output.append(x)
        return x

    def backward(self, p, q):
        # print(p, q)
        self.input.append(q)
        self.output.append(p)
        for i, output in reversed(list(enumerate(self.output))):
            
            if i == (len(self.output) - 1):
                # for last node, use g
                output.backward()
            else:
                output.backward(self.input[i+1].grad.data)
                print(i + 1, self.input[i+1].grad.data)     
        print(i, self.input[i].grad.data)

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

inp = torch.ones(1, 4)
output = model(inp) # predcicted output

output = Variable(output.data, requires_grad=True)
loss = (10 - output).sum()
model.backward(loss, output)


How to use x.register_hook(hook_fn) can you help me with some code snippet?
Thanks!