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

f

Browse Source
This commit is contained in:
Zchen
2025-10-17 01:58:28 +08:00
parent 8ee09b6b5e
commit 49700456b8
1 changed files with 1 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
model_training_nnn_tpu/trainer_tf.py
View File

@@ -98,6 +98,7 @@ class BrainToTextDecoderTrainerTF:
# This ensures we're in the correct replica context # This ensures we're in the correct replica context
print("🔧 Creating optimizer slot variables within TPU replica context...") print("🔧 Creating optimizer slot variables within TPU replica context...")
@tf.function
def init_optimizer_slots(): def init_optimizer_slots():
dummy_gradients = [tf.zeros_like(var) for var in self.model.trainable_variables] dummy_gradients = [tf.zeros_like(var) for var in self.model.trainable_variables]
self.optimizer.apply_gradients(zip(dummy_gradients, self.model.trainable_variables)) self.optimizer.apply_gradients(zip(dummy_gradients, self.model.trainable_variables))