diff --git a/gemma/research/t5gemma/sampling.py b/gemma/research/t5gemma/sampling.py
index 3ea754b2..698b325c 100644
--- a/gemma/research/t5gemma/sampling.py
+++ b/gemma/research/t5gemma/sampling.py
@@ -94,7 +94,7 @@ def tokens(self) -> Int['B L'] | Int['L']:
     """Predicted tokens."""
     return self._maybe_unbatch(self.state.predicted_tokens)
 
-  def _maybe_unbatch(self, x: Array['B *d']) -> Float['*d']:
+  def _maybe_unbatch(self, x: Array['B *d']) -> Array['*d']:
     if isinstance(self.text, str):
       (x,) = x
     return x