Internal change

PiperOrigin-RevId: 461630368
Nitin-Mane · Jul 27, 2022 · 7dcfde4 · 7dcfde4
1 parent 3bb597b
commit 7dcfde4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vit_jax/utils.py b/vit_jax/utils.py
@@ -111,7 +111,7 @@ def acc_grad_and_loss(i, l_and_g):
                                    (step_size, labels.shape[1]))
       li, gi = loss_and_grad_fn(params, imgs, lbls)
       l, g = l_and_g
-      return (l + li, jax.tree_multimap(lambda x, y: x + y, g, gi))
+      return (l + li, jax.tree_map(lambda x, y: x + y, g, gi))
 
     l, g = jax.lax.fori_loop(1, accum_steps, acc_grad_and_loss, (l, g))
     return jax.tree_map(lambda x: x / accum_steps, (l, g))