Input size changed in after ReLU layer in tf.gradient_tape()

Question

I am trying to loop over every layer in the network to obtain the gradient of the network.

When I loop model.layers[layernumber] in tf.gradient_tape(), ReLU layer altered the received input size from (466,36) to (466,195).

However, when I check the summary of my layer it is still (466,36). What is the problem, and why only the size changed in tf.gradient_tape()?

Here is my code:

# Convert the input sample to a tensor
input_sample = X[0]
input_tensor = input_sample.reshape((1,) + input_sample.shape)
input_tensor = tf.convert_to_tensor(input_tensor, dtype=tf.float32)
input_tensor = tf.Variable(input_tensor, dtype=tf.float32, trainable=True)

# Record operations on the tape to calculate gradients
layers = model.layers

# Calculate the gradients for each layer
gradients = []

# tf.compat.v1.disable_eager_execution()
# Iterate through each layer
for layer in layers:
    with tf.GradientTape() as tape:
        tape.watch(input_tensor)

        # Forward pass through the layer
        output = layer(input_tensor)

        print(output.shape, layer)

    # Calculate the gradients of the layer output with respect to the input tensor
    layer_gradients = tape.gradient(output, input_tensor)

    gradients.append(layer_gradients)

I tried to change ReLU into Lamda layer with ReLU but it is still not working. I also changed it to only obtaiin the gradient of the output and input and same problem still occurs`