diff --git a/optformer/omnipred/omnipred.py b/optformer/omnipred/omnipred.py
index 5362af9..dd63b64 100644
--- a/optformer/omnipred/omnipred.py
+++ b/optformer/omnipred/omnipred.py
@@ -108,6 +108,7 @@ def _sample_tokens(
         rng=subkey,
         return_all_decodes=True,
         num_decodes=self.num_samples,
+        prompt_with_targets=False,
     )
     return sampled_tokens, aux['scores']