init

app.py

@@ -12,13 +12,15 @@ import wandb
 
 cfg = load_cfg("configs/effb0-base-breakhis.yaml")
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 model = system.SimpleClassificationSystem.load_from_checkpoint(
     "runs/03_20_2025_20_19_28/wandb/histopath/03_20_2025_20_19_28/checkpoints/epoch=210-step=41778.ckpt",
     torch.device("cpu"),
     cfg=cfg.system
 )
 
-model.to("cuda")
+model.to(device)
 model.eval()
 
 print("Model loaded successfully!")
@@ -34,7 +36,7 @@ def preprocess_image(image: Image):
         v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
     ])
     
-    return transform(image).unsqueeze(0).to("cuda")
+    return transform(image).unsqueeze(0).to(device)
 
 
 def predict(image_path: str):