Update app.py

app.py

@@ -1,17 +1,23 @@
 import torch
-from altair import Description
 from sentence_transformers import SentenceTransformer
 from datasets import load_dataset
 from sentence_transformers.util import cos_sim
 import gradio as gr
 
+# --- FIX: Dynamically select the device ---
 if torch.backends.mps.is_available():
-    print("MPS is available. Using M1 GPU!")
+    device = torch.device("mps")
+    print("MPS device is available. Using M1/M2 GPU!")
+elif torch.cuda.is_available():
+    device = torch.device("cuda")
+    print("CUDA device is available. Using NVIDIA GPU!")
 else:
-    print("MPS not available. Check your PyTorch installation.")
+    device = torch.device("cpu")
+    print("No GPU available. Falling back to CPU.")
+
 
 model_name = 'all-MiniLM-L6-v2'
-model = SentenceTransformer(model_name, device='mps') # Use 'mps' to utilize the M1 chip
+model = SentenceTransformer(model_name, device=device)
 print("Loading dataset...")
 
 ds = load_dataset("pszemraj/goodreads-bookgenres", "default")
@@ -62,4 +68,4 @@ gr_interface = gr.Interface(
     examples=["A thrilling detective story", "A heartwarming novel about friendship", "Science fiction about space travel"]
 )
 
-gr_interface.launch(share=True)
+gr_interface.launch(share=True)