From 6e11097f06a4b9e81e36e1ac004bfef953c018c2 Mon Sep 17 00:00:00 2001
From: Xingyou Song <xingyousong@google.com>
Date: Sun, 17 Mar 2024 08:48:56 -0700
Subject: [PATCH] Internal change

PiperOrigin-RevId: 616604781
---
 optformer/omnipred/omnipred.py | 1 +
 1 file changed, 1 insertion(+)

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']