Merge branch 'mllm' of github.com:mlfoundations/open_flamingo into im…

…age_res
mlfoundations · Nov 30, 2023 · 2053d1f · 2053d1f
2 parents adc61a6 + c5feb97
commit 2053d1f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/open_flamingo/train/losses.py b/open_flamingo/train/losses.py
@@ -88,7 +88,7 @@ def __call__(
         shift_labels = labels[..., 1:].contiguous()
         # Flatten the tokens
         loss_fct = CrossEntropyLossWithZLoss(eps=z_loss_eps)
-        shift_logits = shift_logits.view(-1, model.lang_model.config.vocab_size)
+        shift_logits = shift_logits.view(-1, unwrap_model(model).lang_model.config.vocab_size)
         shift_labels = shift_labels.view(-1)
         # Enable model parallelism
         shift_labels = shift_labels.to(shift_logits.device)