I am working on the following model, and I want to calculate the gradient of input to each layer; pseudo-code-
class MyModel():
    def __init__():
        layers ....
        ....
        ..
    def forward(x):
        layer_input = []
        layer_input.append(x)
        for layer in layers:
            x = layer(x)
            layer_input.append(x)

        return x, layer_input


and I want to do this, also so how to freeze the weights and not an input_to_each layer?
for epoch in epochs:
    #train model
    for images, labels in data:
        out, layer_input = model(images)
        ...

    # validate model
    for images, labels in data:
        out, layer_input = model(images)
        ...

    # save gradient for eacch layer input
    # also freeze the layers and not a input to each layer (??)
    for images, labels in data:
        out, layer_input = model(images)
        loss = criterion(out, labels)
        for l_input in layer_input:
            grad = calculate_grad(loss, l_input)
            save(grad)

Thank you in advance.