Update main.py

main.py

@@ -185,6 +185,9 @@ async def lifespan(app: FastAPI):
         if os.path.exists("DECODE_Frequency_Twin.pth"):
             ckpt = torch.load("DECODE_Frequency_Twin.pth", map_location='cpu')
             sd = ckpt['model_state_dict'] if isinstance(ckpt, dict) and 'model_state_dict' in ckpt else ckpt
+            # Remove unexpected 'theta_h' key if present
+            if 'theta_h' in sd:
+                sd = {k: v for k, v in sd.items() if k != 'theta_h'}
             model = FrequencyPINN()
             model.load_state_dict(sd, strict=True)
             ml_assets["f_model"] = model.eval()