Update cross_encoder_reranking_train.py

cross_encoder_reranking_train.py

@@ -11,11 +11,14 @@ from transformers import AutoTokenizer, AutoModel
 from sklearn.cluster import AgglomerativeClustering
 from sklearn.metrics.pairwise import cosine_similarity
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 # Load embedder once
 embedder = SentenceTransformer("all-MiniLM-L6-v2")
+embedder = embedder.to(device)
+
 
 def embed_text_list(texts):
-    return embedder.encode(texts, convert_to_tensor=False)
+    return embedder.encode(texts, convert_to_tensor=False, device=device)
 
 def rank_by_centrality(texts):
     embeddings = embed_text_list(texts)