Make ClipSeg compatible with model parallelism (#22844)

84a6570e · Youssef Adarrab · GitHub · 5bb4ec62 · 84a6570e
Unverified Commit 84a6570e authored Apr 18, 2023 by Youssef Adarrab Committed by GitHub Apr 18, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 0 deletions

src/transformers/models/clipseg/modeling_clipseg.py src/transformers/models/clipseg/modeling_clipseg.py +2 -0

No files found.
--- a/src/transformers/models/clipseg/modeling_clipseg.py
+++ b/src/transformers/models/clipseg/modeling_clipseg.py
@@ -1480,6 +1480,8 @@ class CLIPSegForImageSegmentation(CLIPSegPreTrainedModel):

        loss = None
        if labels is not None:
+            # move labels to the correct device to enable PP
+            labels = labels.to(logits.device)
            loss_fn = nn.BCEWithLogitsLoss()
            loss = loss_fn(logits, labels)