WIP

1 files changed, 2 insertions, 2 deletions

diff --git a/finetuning.py b/finetuning.py
index 853eb81..62f506c 100755
--- a/finetuning.py
+++ b/finetuning.py
@@ -32,8 +32,8 @@ def add_chat_text(example, processor):
     return example
 
 
-def collate(batch, processor, max_length):  # FIXME: Support batch_size > 1
-    images  = [i["image"] for i in batch]
+def collate(batch, processor, max_length):
+    images  = [[i["image"]] for i in batch]
     texts = [i["text"]  for i in batch]
     processor_output = processor(
         text=texts,