No public description

PiperOrigin-RevId: 568331404
google-research · Sep 25, 2023 · 75af72b · 75af72b
1 parent dd822ef
commit 75af72b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vmoe/evaluate/evaluator.py b/vmoe/evaluate/evaluator.py
@@ -160,7 +160,7 @@ def _make_callback_fn(cls, *, apply_fn, loss_fn, label_pred_fn, datasets,
         in_shardings=(), out_shardings=None, static_argnums=(0,))
     def make_eval_state_pjit(seed):
       rngs = utils.make_rngs(rng_keys, seed)
-      return EvalState(
+      return EvalState(  # pytype: disable=wrong-arg-types  # jnp-type
           num=jnp.zeros((), dtype=jnp.float32),
           sum_correct=jnp.zeros((), dtype=jnp.float32),
           sum_loss=jnp.zeros((), dtype=jnp.float32),