From 49700456b8dd029d6bb3ac33ff8e3a67f8f56584 Mon Sep 17 00:00:00 2001 From: Zchen <161216199+ZH-CEN@users.noreply.github.com> Date: Fri, 17 Oct 2025 01:58:28 +0800 Subject: [PATCH] f --- model_training_nnn_tpu/trainer_tf.py | 1 + 1 file changed, 1 insertion(+) diff --git a/model_training_nnn_tpu/trainer_tf.py b/model_training_nnn_tpu/trainer_tf.py index e59b2f0..0a72551 100644 --- a/model_training_nnn_tpu/trainer_tf.py +++ b/model_training_nnn_tpu/trainer_tf.py @@ -98,6 +98,7 @@ class BrainToTextDecoderTrainerTF: # This ensures we're in the correct replica context print("🔧 Creating optimizer slot variables within TPU replica context...") + @tf.function def init_optimizer_slots(): dummy_gradients = [tf.zeros_like(var) for var in self.model.trainable_variables] self.optimizer.apply_gradients(zip(dummy_gradients, self.model.trainable_variables))