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")