Fix device_id in mmc4 pass (#250)

mlfoundations · Aug 28, 2023 · a05dcba · a05dcba
1 parent 9bebf9f
commit a05dcba
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/open_flamingo/train/train_utils.py b/open_flamingo/train/train_utils.py
@@ -155,8 +155,8 @@ def train_one_epoch(
         with autocast():
             loss_mmc4 = model(
                 vision_x=images,
-                lang_x=input_ids,
-                attention_mask=attention_mask,
+                lang_x=input_ids.to(device_id),
+                attention_mask=attention_mask.to(device_id),
                 labels=labels,
             )[0]