Update app.py

app.py

@@ -66,18 +66,22 @@ class SleepStageCNN(nn.Module):
 # ────────────────────────────────────────────────────────────────
 # Load Model at startup
 # ────────────────────────────────────────────────────────────────
-
 device = torch.device("cpu")
 model = SleepStageCNN(n_channels=1, n_classes=6)
-
 if os.path.exists(MODEL_PATH):
     checkpoint = torch.load(
         MODEL_PATH, map_location=device, weights_only=False
     )
     if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
-        model.load_state_dict(checkpoint["model_state_dict"])
+        state_dict = checkpoint["model_state_dict"]
     else:
-        model.load_state_dict(checkpoint)
+        state_dict = checkpoint
+
+    # Remap bare Sequential keys (e.g. "0.weight") → "network.0.weight"
+    if any(k.split(".")[0].isdigit() for k in state_dict.keys()):
+        state_dict = {"network." + k: v for k, v in state_dict.items()}
+
+    model.load_state_dict(state_dict)
     model.eval().to(device)
     print(f"✅ Model loaded from {MODEL_PATH}")
 else:
@@ -85,8 +89,6 @@ else:
         f"Model file not found at {MODEL_PATH}. "
         "Upload sleep_stage_cnn.pth to this Space."
     )
-
-
 # ────────────────────────────────────────────────────────────────
 # Inference Function
 # ────────────────────────────────────────────────────────────────