From a0b59c6987effafe2b75b976453f359ed6383b59 Mon Sep 17 00:00:00 2001
From: Zchen <161216199+ZH-CEN@users.noreply.github.com>
Date: Thu, 16 Oct 2025 21:06:01 +0800
Subject: [PATCH] fix

---
 model_training_nnn_tpu/trainer_tf.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/model_training_nnn_tpu/trainer_tf.py b/model_training_nnn_tpu/trainer_tf.py
index d5dfb91..4183208 100644
--- a/model_training_nnn_tpu/trainer_tf.py
+++ b/model_training_nnn_tpu/trainer_tf.py
@@ -92,6 +92,13 @@ class BrainToTextDecoderTrainerTF:
             self.optimizer = self._create_optimizer()
             print("✅ Optimizer created")
 
+            print("🔧 Pre-building optimizer state for TPU...")
+            # Force optimizer to build its internal state within strategy scope
+            # This prevents the 'NoneType' strategy error during first apply_gradients
+            dummy_grads = [tf.zeros_like(w) for w in self.model.trainable_variables]
+            self.optimizer.apply_gradients(zip(dummy_grads, self.model.trainable_variables))
+            print("✅ Optimizer state pre-built successfully")
+
             print("📅 Setting up learning rate scheduler...")
             self.lr_scheduler = self._create_lr_scheduler()
             print("✅ LR scheduler ready")