Норма партии Tensorflow 2 не вызывает градиента

Question

for batch_step, batch in enumerate(self.batch_iterator_train):
    with tf.GradientTape() as tape:
         logits, _ = self.model(batch['images'], is_training=True)
         loss_value = self.loss_fn(labels=batch['labels'], logits=logits)
    grads = tape.gradient(loss_value, self.model.trainable_variables) # use gradient to automatically retrieve gradients of var. wrt loss
    self.optimizer.apply_gradients(zip(grads, self.model.trainable_variables)) # apply one step of gradient descent

Я использую приведенный выше код для обучения моей модели в Tensorflow 2.1.0, но получаю следующую ошибку:

WARNING:tensorflow:Gradients do not exist for variables ['net/block/batch_norm_relu_1/batch_normalization_1/beta:0', .....]

Есть ли проблема с TF 2.10 с пакетом норма, особенно с бета-параметром? Насколько я понимаю, бета-параметр используется для смещения значений нормы партии, поэтому он должен быть обучаемым.

ОБНОВЛЕНИЕ: При более внимательном рассмотрении я заметил, что около половины моих слоев нормы партии имеют веса всех 0s.

self.model.trainable_variables[17]
Out[63]: 
array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
      dtype=float32)>

Хотя я действительно не уверен, почему.