move to(device) outside of collate

utils.py

@@ -226,7 +226,7 @@ def collate_fn(examples, tokenizer=None, padding=None, device=None):
             batch[k].append(v)
 
     return {
-        k: torch.tensor(v, dtype=torch.long, device=device) if k in {"attention_mask", "input_ids"} else v for k, v in batch.items()
+        k: torch.tensor(v, dtype=torch.long) if k in {"attention_mask", "input_ids"} else v for k, v in batch.items()
     }
 
 @torch.inference_mode()
@@ -332,8 +332,8 @@ def batch_embed(
         drop_last=False,
         collate_fn=partial(collate_fn, device=device)
     ):
-        ids = batch["input_ids"]
-        mask = batch["attention_mask"]
+        ids = batch["input_ids"].to(device)
+        mask = batch["attention_mask"].to(device)
 
         t_ids = torch.zeros_like(ids)