1 files changed, 1 insertions, 1 deletions

diff --git a/data/csv.py b/data/csv.py
index df3ee77..b058a3e 100644
--- a/data/csv.py
+++ b/data/csv.py
@@ -121,7 +121,7 @@ def collate_fn(weight_dtype: torch.dtype, tokenizer: CLIPTokenizer, examples):
     inputs = unify_input_ids(tokenizer, input_ids)
 
     batch = {
-        "with_prior": torch.tensor(with_prior),
+        "with_prior": torch.tensor([with_prior] * len(examples)),
         "prompt_ids": prompts.input_ids,
         "nprompt_ids": nprompts.input_ids,
         "input_ids": inputs.input_ids,