Update collator.py

Former-commit-id: 941fa8a0d9c3a9106ad0af6e776db7e57f69548f
This commit is contained in:
hoshi-hiyouga 2024-10-29 22:03:42 +08:00 committed by GitHub
parent a9afffa246
commit 58fb24ce41

View File

@ -99,8 +99,8 @@ class MultiModalDataCollatorForSeq2Seq(DataCollatorForSeq2Seq):
features: Dict[str, "torch.Tensor"] = super().__call__(features) features: Dict[str, "torch.Tensor"] = super().__call__(features)
features.update(mm_inputs) features.update(mm_inputs)
if features.get("pixel_values") is not None and isinstance(features["pixel_values"], list): if isinstance(features.get("pixel_values"), list): # for pixtral inputs
features = features.data features = features.data # use default_collate() instead of BatchEncoding.to()
return features return features