This commit is contained in:
Zchen
2025-10-16 21:06:01 +08:00
parent ed6e21bfe4
commit a0b59c6987

View File

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