I have an MLP model and I want to save the gradient after each iteration and average it at the last. How I can do that?
model:
class MyModel(torch.nn.Module):

    def __init__(self, layers_size, input_size=784, num_classes=10):

        super(MyModel, self).__init__()

        # create input layer
        self.layers = torch.nn.ModuleList(
                    [torch.nn.Linear(input_size, layers_size[0])])
        
        # stack hidden layer
        self.layers.extend([torch.nn.Linear(layers_size[i - 1], 
                        layers_size[i]) for i in range(1, len(layers_size ))])
        
        # add output layer
        self.layers.append(torch.nn.Linear(layers_size[-1], num_classes))
        
    def forward(self, x):
        # iterate over all the layers
        for i, layer in enumerate(self.layers):
            if i == len(self.layers):
                x = func.softmax(layer(x)) # softmax for last layer
            else:
                x = func.relu(layer(x))
        return x


Here what I like to do:
for epoch in epochs
	for batch in batches:
		model.forward(batch)
		compute_gradients;
		save(gradients)
		model.backward()

avarage(gradients)

Thanks in advance.