Fixed cpu mapping error

app.py

@@ -70,12 +70,9 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
 # Load YOLO detector
-#yolo = YOLO("/home/sankalp/flake_classification/models/best.pt")
-#yolo = YOLO("/home/sankalp/yolo_flake_detection/yolo11n_synthetic_runs/exp1/weights/best.pt")
-#yolo = YOLO("/home/sankalp/yolo_flake_detection/yolo_runs/yolo11l_flake_runs/weights/best.pt")
-
 yolo = YOLO("models/uark_detector_v3.pt")
-torch.load("models/flake_classifier_5layer.pth",map_location=torch.device("cpu"))
+
+# Load classifier model checkpoint with proper device mapping
 ckpt_path = "models/flake_classifier_5layer.pth"
 ckpt = torch.load(ckpt_path, map_location=device)
 
@@ -139,4 +136,4 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(share=True)