# Tape.gradient yields wrong dimension (when dynamically freezing layers)

**URL:** <https://discuss.ai.google.dev/t/tape-gradient-yields-wrong-dimension-when-dynamically-freezing-layers/109648>\
**Category:** TensorFlow\
**Tags:** tfgradient\
**Created:** [November 17, 2025, 11:08am UTC](https://discuss.ai.google.dev/t/tape-gradient-yields-wrong-dimension-when-dynamically-freezing-layers/109648 "2025-11-17T11:08:55Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![Cola\_Lightyear](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/cola_lightyear/32/15884_2.png) [@Cola\_Lightyear](https://discuss.ai.google.dev/u/Cola_Lightyear)\
**Post date:** [November 17, 2025, 11:08am UTC](https://discuss.ai.google.dev/t/tape-gradient-yields-wrong-dimension-when-dynamically-freezing-layers/109648/1 "2025-11-17T11:08:56Z")

</div>

Hi, I am trying to train a model that consists of certain submodels A,B, C and D.  
To get a stable training, I am trying to iteratively train these submodels by freezing all of them except one, train the unfrozen submodel for a few iterations and to then switch the frozen submodel to another.

I use gradient.tape to track the gradients of the trainable variables. However, the

```auto
gradients_total = tape.gradient(total_loss, model.trainable_variables)

```

call yields a list of gradients that does not match the number of actually trainable variables. Then the apply\_gradients step fails.

Here is the gradient function:  
`@tf.function`  
`def compute_gradients(x_batch, y_batch, model, loss_fn, use_sisdr_loss):`  
`with tf.GradientTape() as tape:`  
`y_pred = model(x_batch, training=True)`

```auto
    loss_sisdr = loss_fn(y_batch[0], y_pred[0]) + loss_fn(y_batch[1], y_pred[1]) # Verlustberechnung
    loss_mse_and_l2 = tf.add_n(model.losses)
    total_loss = loss_sisdr + loss_mse_and_l2

if use_sisdr_loss:
    gradients_total = tape.gradient(total_loss, model.trainable_variables)
else:
    gradients_total = tape.gradient(loss_mse_and_l2, model.trainable_variables)

return gradients_total, total_loss, loss_sisdr

```

`def train_step(x_batch, y_batch, model, loss_fn, optimizer, use_sisdr_loss, compute_grads, all_trainable_variables):`  
`global buffer, window_size, min_window_size, max_window_size`  
`gradients_total, loss, loss_sisdr = compute_grads(x_batch, y_batch, model, loss_fn, use_sisdr_loss)`

I suspect that this behavior is due to the @tf.function decorator. I use a callback during training to freeze/unfreeze the submodels. Here is how I do that:  
`class ModelCircularFreezeCallback(tf.keras.callbacks.Callback):`  
`def ` **`init`** `(self, init_active_model=0, switch_iter=5):`  
`super().` **`init`** `()`  
`self.active_model_idx = init_active_model`  
`self.switch_iter = switch_iter`

```auto
def SetActiveModel(self):
    print(f"Unfreezing RNN {self.active_model_idx} - freezing all other RNNs")
    # frozen_vars = []
    for k,ssrae_name in enumerate(ssrae_names):
        l = self.model.get_layer(ssrae_name)
        if k == self.active_model_idx:
            l.trainable = True
            print(f"Unfreezed {ssrae_name}")
            for i in range(len(l.layers)):
                l.layers[i].trainable = True
        else:
            # vars_ssrae = list(l.variables)
            l.trainable = False
            for i in range(len(l.layers)):
                l.layers[i].trainable = False`

```

I cannot get rid of the @tf.function decorator, because the training is way too slow without it and I run out of memory.  
Does anyone know how I can dynamically freeze/unfreeze layers during training?  
Thank you!

---

<div class="post-metadata">

**Author:** ![Cola\_Lightyear](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/cola_lightyear/32/15884_2.png) [@Cola\_Lightyear](https://discuss.ai.google.dev/u/Cola_Lightyear)\
**Post date:** [November 17, 2025, 2:28pm UTC](https://discuss.ai.google.dev/t/tape-gradient-yields-wrong-dimension-when-dynamically-freezing-layers/109648/3 "2025-11-17T14:28:06Z")

</div>

As a workaround I introduced a mask to the gradient application where gradients of the weights of frozen layers are now substituted with “None”.  
This appears to work, but a better approach is still welcome!

---

<div class="post-metadata">

**Author:** ![Divya\_Sree\_Kayyuri](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/divya_sree_kayyuri/32/29791_2.png) [@Divya\_Sree\_Kayyuri](https://discuss.ai.google.dev/u/Divya_Sree_Kayyuri)\
**Post date:** [February 11, 2026, 4:52am UTC](https://discuss.ai.google.dev/t/tape-gradient-yields-wrong-dimension-when-dynamically-freezing-layers/109648/6 "2026-02-11T04:52:16Z")

</div>

Hi @Cola_Lightyear, When `@tf.function` is used, `model.trainable_variables` is traced only once during the initial tracing of the function. To fix this, pass `model.trainable_variables` as an explicit argument to your `@tf.function`. Thanks!