add optional cuda

colpali_engine/trainer/retrieval_evaluator.py

@@ -1,11 +1,12 @@
 import torch
 from mteb.evaluation.evaluators import RetrievalEvaluator
-
+from colpali_engine.utils.torch_utils import get_torch_device
 
 class CustomEvaluator:
     def __init__(self, is_multi_vector=False):
         self.is_multi_vector = is_multi_vector
         self.mteb_evaluator = RetrievalEvaluator()
+        self.device = get_torch_device()
 
     def evaluate(self, qs, ps):
         if self.is_multi_vector:
@@ -51,12 +52,12 @@ def evaluate_colbert(self, qs, ps, batch_size=128) -> torch.Tensor:
         for i in range(0, len(qs), batch_size):
             scores_batch = []
             qs_batch = torch.nn.utils.rnn.pad_sequence(qs[i : i + batch_size], batch_first=True, padding_value=0).to(
-                "cuda"
+                self.device
             )
             for j in range(0, len(ps), batch_size):
                 ps_batch = torch.nn.utils.rnn.pad_sequence(
                     ps[j : j + batch_size], batch_first=True, padding_value=0
-                ).to("cuda")
+                ).to(self.device)
                 scores_batch.append(torch.einsum("bnd,csd->bcns", qs_batch, ps_batch).max(dim=3)[0].sum(dim=2))
             scores_batch = torch.cat(scores_batch, dim=1).cpu()
             scores.append(scores_batch)