Ignore incorrect type annotations related to jax dtypes

PiperOrigin-RevId: 571882947
google · Oct 9, 2023 · 34825d3 · 34825d3
1 parent ca9ec51
commit 34825d3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/flax_integration/main.py b/examples/flax_integration/main.py
@@ -151,7 +151,7 @@ def _loss_fn(params):
       scores = model.apply(
           flax.core.copy(model_state, {"params": params}), inputs
       )
-      loss = loss_fn(scores, labels, where=mask, reduce_fn=jnp.mean)
+      loss = loss_fn(scores, labels, where=mask, reduce_fn=jnp.mean)  # pytype: disable=wrong-arg-types  # jnp-type
       return loss
 
     params = model_state["params"]