fix: mllama missing aspect ratio

2025-07-22 10:36:07 +07:00
parent 78169cd37a
commit ee19007ba4
1 changed files with 3 additions and 0 deletions
--- a/src/axolotl/utils/collators/mm_chat.py
+++ b/src/axolotl/utils/collators/mm_chat.py
@@ -92,6 +92,9 @@ class MultiModalChatDataCollator(DataCollatorMixin):
        if "pixel_values" in batch:
            final_batch["pixel_values"] = torch.stack(batch["pixel_values"])

+        if "aspect_ratio_ids" in batch:
+            final_batch["aspect_ratio_ids"] = torch.stack(batch["aspect_ratio_ids"])
+
        if "input_features" in batch:
            final_batch["input_features"] = torch.stack(batch["input_features"])