diff --git a/gemma/gm/text/_chat_sampler.py b/gemma/gm/text/_chat_sampler.py
index 05710159..b230da6f 100644
--- a/gemma/gm/text/_chat_sampler.py
+++ b/gemma/gm/text/_chat_sampler.py
@@ -256,6 +256,7 @@ def _sample(
           return_state=True,
           last_state=last_state,
           stream=bool(stream),
+          sharding=sharding,
       )
 
   def chat(