zchen/b2txt25
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
Zchen
2025-10-16 21:06:01 +08:00
parent ed6e21bfe4
commit a0b59c6987
1 changed files with 7 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
model_training_nnn_tpu/trainer_tf.py
View File

@@ -92,6 +92,13 @@ class BrainToTextDecoderTrainerTF:
self.optimizer = self._create_optimizer() self.optimizer = self._create_optimizer()
print("✅ Optimizer created") 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...") print("📅 Setting up learning rate scheduler...")
self.lr_scheduler = self._create_lr_scheduler() self.lr_scheduler = self._create_lr_scheduler()
print("✅ LR scheduler ready") print("✅ LR scheduler ready")