fix collator

utils.py

@@ -2,6 +2,7 @@ import os
 import time
 import shutil
 from pathlib import Path
+from functools import partial
 from typing import Union, Dict, List
 
 import torch
@@ -321,6 +322,7 @@ def batch_embed(
 
     start_time = time.time()
 
+    
     for batch in DataLoader(
         ds,
         batch_size=inference_bs,
@@ -328,8 +330,8 @@ def batch_embed(
         num_workers=1,
         pin_memory=True,
         drop_last=False,
+        collate_fn=partial(collate_fn, device=device)
     ):
-        batch = collate_fn(batch, device=device)
         ids = batch["input_ids"]
         mask = batch["attention_mask"]