Update main.py

main.py

@@ -171,6 +171,10 @@ async def lifespan(app: FastAPI):
         if os.path.exists("battery_model.pt"):
             ckpt = torch.load("battery_model.pt", map_location='cpu')
             sd = ckpt['model_state_dict'] if isinstance(ckpt, dict) and 'model_state_dict' in ckpt else ckpt
+            # Remove unexpected 'raw_R' and 'raw_h' keys if present
+            for k in ["raw_R", "raw_h"]:
+                if k in sd:
+                    sd = {key: val for key, val in sd.items() if key != k}
             model = BatteryPINN()
             model.load_state_dict(sd, strict=True)
             ml_assets["b_model"] = model.eval()