Update app.py

app.py

@@ -8,6 +8,8 @@ model = AutoModelForSequenceClassification.from_pretrained(
     trust_remote_code=True,
 )
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+model.to(device)  # Move model to GPU if available, otherwise CPU
 model.eval()
 
 def compute_scores(query, documents):
@@ -24,7 +26,7 @@ def compute_scores(query, documents):
     documents_list = documents.split('\n')
     sentence_pairs = [[query, doc] for doc in documents_list]
     scores = model.compute_score(sentence_pairs, max_length=1024)
-    return scores.tolist()
+    return scores
 
 # Define Gradio interface
 iface = gr.Interface(
@@ -40,4 +42,4 @@ iface = gr.Interface(
 
 # Launch the interface
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()