Skip to content

GitLab

Unverified Commit e5c12c03 authored by Bram Willemsen's avatar Bram Willemsen Committed by GitHub
Browse files

Move tensors to same device to enable IDEFICS naive MP training (#27746)

parent 3e68944c
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
src/transformers/models/idefics/modeling_idefics.py
View file @ e5c12c03
...@@ -1510,9 +1510,10 @@ class IdeficsForVisionText2Text(IdeficsPreTrainedModel): ...@@ -1510,9 +1510,10 @@ class IdeficsForVisionText2Text(IdeficsPreTrainedModel):
loss = None loss = None
if labels is not None: if labels is not None:
labels = labels.to(logits.device)
# Shift so that tokens < n predict n # Shift so that tokens < n predict n
if attention_mask is not None: if attention_mask is not None:
shift_attention_mask = attention_mask[..., 1:] shift_attention_mask = attention_mask[..., 1:].to(logits.device)
shift_logits = logits[..., :-1, :][shift_attention_mask != 0].contiguous() shift_logits = logits[..., :-1, :][shift_attention_mask != 0].contiguous()
shift_labels = labels[..., 1:][shift_attention_mask != 0].contiguous() shift_labels = labels[..., 1:][shift_attention_mask != 0].contiguous()
else: else:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment