fix batching

utils.py

@@ -312,14 +312,14 @@ def batch_embed(
         pin_memory=True,
         drop_last=False,
     ):
-        ids = torch.tensor([b["input_ids"] for b in batch], device=device)
-        mask = torch.tensor([b["attention_mask"] for b in batch], device=device)
+        ids = torch.tensor(batch["input_ids"], device=device)
+        mask = torch.tensor(batch["attention_mask"], device=device)
         t_ids = torch.zeros_like(ids)
 
         outputs = model(input_ids=ids, attention_mask=mask, token_type_ids=t_ids)
 
         embeds.extend(mean_pooling(outputs[0], mask).cpu().tolist())
-        texts.extend([b[column_name] for b in batch])
+        texts.extend(batch[column_name])
 
         current_count += ids.shape[0]