Spaces:

sanpdy
/

quantum-flake-pipeline

Sleeping

sanpdy commited on Jul 15

Commit

3d424a0

1 Parent(s): 2057317

Fixed cpu mapping error

Files changed (1) hide show

app.py CHANGED Viewed

@@ -75,7 +75,7 @@ print(f"Using device: {device}")
 #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")
 ckpt_path = "models/flake_classifier_5layer.pth"
 ckpt = torch.load(ckpt_path, map_location=device)

 #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"))
 ckpt_path = "models/flake_classifier_5layer.pth"
 ckpt = torch.load(ckpt_path, map_location=device)