Update app.py

app.py

@@ -11,15 +11,22 @@ import gradio as gr
 
 import models
 
+print(f"Is CUDA available: {torch.cuda.is_available()}")
+# True
+print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+# Tesla T4
+
 with open("index_to_species.json", "r") as file:
     index_to_species_data = file.read()
 index_to_species = json.loads(index_to_species_data)
 
 num_classes = len(list(index_to_species.keys()))
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 # Load the model
 classify_model = models.DinoVisionTransformerClassifier(num_classes)
-classify_model.load_state_dict(torch.load("best_dinov2_both_2023-11-21_07-44-35.pth", map_location=torch.device('cpu')))
+classify_model.load_state_dict(torch.load("best_dinov2_both_2023-11-21_07-44-35.pth", map_location=torch.device(device)))
 classify_model.eval()
 
 k = 5