fix cpu error

handler.py

@@ -1,6 +1,7 @@
 from pylate import models
 from transformers import AutoTokenizer
 import torch
+import numpy as np
 
 class EndpointHandler:
     def __init__(self, path=""):
@@ -8,6 +9,21 @@ class EndpointHandler:
         self.model = models.ColBERT(model_name_or_path=path)
         self.model.eval()
 
+    def _to_list(self, emb):
+        """
+        Make the output JSON-serialisable:
+        – torch.Tensor  ➜  emb.cpu().tolist()
+        – np.ndarray    ➜  emb.tolist()
+        – list[...]     ➜  recurse
+        """
+        if isinstance(emb, torch.Tensor):
+            return emb.cpu().tolist()
+        if isinstance(emb, np.ndarray):
+            return emb.tolist()
+        if isinstance(emb, list):
+            return [self._to_list(e) for e in emb]
+        return emb  # already plain Python
+
     def __call__(self, data):
         texts = data.get("inputs") or data.get("text") or data
         if isinstance(texts, str):
@@ -16,8 +32,8 @@ class EndpointHandler:
         with torch.no_grad():
             emb = self.model.encode(
                 texts,
-                is_query=True,          # query-style encoding
+                is_query=True,
                 batch_size=32,
             )
-        # TEI expects JSON-serialisable output
-        return emb.cpu().tolist()
+
+        return self._to_list(emb)