Update src/streamlit_app.py

src/streamlit_app.py

@@ -18,7 +18,16 @@ IMAGE_SIZE = 300
 # =========================
 @st.cache_resource
 def load_model():
-    model = torch.load(MODEL_PATH, map_location="cpu")
+    MODEL_PATH = Path(__file__).resolve().parent / "efficientnet_b3_best.pt"
+
+    model = timm.create_model(
+        "efficientnet_b3",
+        pretrained=False,
+        num_classes=2
+    )
+
+    state_dict = torch.load(MODEL_PATH, map_location="cpu")
+    model.load_state_dict(state_dict)
     model.eval()
     return model