Minor bug fix in train step.

PiperOrigin-RevId: 558731486
google-research · Aug 21, 2023 · dd822ef · dd822ef
1 parent 760d672
commit dd822ef
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vmoe/train/trainer.py b/vmoe/train/trainer.py
@@ -642,7 +642,7 @@ def train_step(
 
   @functools.partial(jax.grad, has_aux=True)
   def compute_grads_and_metrics(params, images, labels, rngs):
-    rngs, next_rngs = utils.tree_rngs_split(state.rngs)
+    rngs, next_rngs = utils.tree_rngs_split(rngs)
     logits, metrics = state.apply_fn({'params': params}, images, rngs=rngs)
     metrics = dict(**metrics)
     metrics['main_loss'] = jnp.mean(loss_fn(logits, labels))