WIP

1 files changed, 1 insertions, 1 deletions

diff --git a/make_dataset.py b/make_dataset.py
index f996acc..65338d5 100755
--- a/make_dataset.py
+++ b/make_dataset.py
@@ -26,7 +26,7 @@ def make_dataset(base="./baseline"):  # TODO: Make actual hf dataset
         }, ensure_ascii=False, indent=0))
         images.append(image_path)
 
-    return Dataset.from_dict({"image": images, "user": user_prompts, "assistant": assistant_replies}).cast_column("image", Image())
+    return Dataset.from_dict({"image": images, "user": user_prompts, "assistant": assistant_replies}).cast_column("image", Image().convert("RGB"))
 
 
 def main():