update rng dtype in predict_batch

google-research · Mar 8, 2024 · 06be7c2 · 06be7c2
1 parent 189868b
commit 06be7c2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/t5x/models.py b/t5x/models.py
@@ -638,7 +638,7 @@ def _compute_kv_cache(
   def predict_batch(self,
                     params: PyTreeDef,
                     batch: Mapping[str, jnp.ndarray],
-                    rng: Optional[jax.random.KeyArray] = None,
+                    rng: Optional[jax.Array] = None,
                     flax_mutables: Optional[PyTreeDef] = None) -> jnp.ndarray:
     """Predicts a batch of outputs from the model.