tiny fix

Former-commit-id: 830511a6d0216da99520aee8b3a753d347a71fa9
2024-08-30 03:21:50 +08:00
parent 30d97ca879
commit 665db18661
8 changed files with 24 additions and 13 deletions
--- a/src/llamafactory/data/collator.py
+++ b/src/llamafactory/data/collator.py
@@ -68,7 +68,7 @@ class CustomDataCollatorForSeq2Seq(DataCollatorForSeq2Seq):
    """

    def __call__(self, features: Sequence[Dict[str, Any]]) -> Dict[str, "torch.Tensor"]:
-        image_grid_thw = None
+        image_grid_thw = None  # TODO: better handle various VLMs
        if "image_grid_thw" in features[0]:
            image_grid_thw_list = [
                torch.Tensor(feature["image_grid_thw"]).long()