Update app.py

app.py

@@ -112,7 +112,7 @@ MAX_LEN = int(CFG.get("max_len", 128))
 
 # ---- Image model
 img_model = timm.create_model(IMG_BACKBONE, pretrained=False, num_classes=NUM_CLASSES)
-sd_img = clean_state_dict(torch.load("best_scin_image.pt", map_location="cpu"))
+sd_img = clean_state_dict(safe_torch_load("best_scin_image.pt"))
 img_model.load_state_dict(sd_img, strict=True)
 img_model.to(DEVICE).eval()
 
@@ -134,7 +134,7 @@ class TextClassifier(nn.Module):
 
 tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_NAME)
 text_model = TextClassifier(TEXT_MODEL_NAME, NUM_CLASSES)
-sd_txt = clean_state_dict(torch.load("best_scin_text.pt", map_location="cpu"))
+sd_txt = clean_state_dict(safe_torch_load("best_scin_text.pt"))
 text_model.load_state_dict(sd_txt, strict=False)
 text_model.to(DEVICE).eval()