Update networks.py

networks.py

@@ -380,7 +380,7 @@ def load_checkpoint(model, checkpoint_path, strict=True):
     # Create a new state dict that matches our model architecture
     new_state_dict = {}
     for key, value in state_dict.items():
-        # Handle any name changes here if needed
+        # Handle name changes
         new_key = key
         if 'gridGen' in key:
             # Map old parameter names to new ones
@@ -388,6 +388,10 @@ def load_checkpoint(model, checkpoint_path, strict=True):
                 new_key = key.replace('P_X', 'P_X_base')
             elif 'P_Y' in key and 'base' not in key:
                 new_key = key.replace('P_Y', 'P_Y_base')
+            elif 'grid_X' in key and 'base' not in key:
+                new_key = key.replace('grid_X', 'grid_X_base')
+            elif 'grid_Y' in key and 'base' not in key:
+                new_key = key.replace('grid_Y', 'grid_Y_base')
         
         # Only include keys that exist in the current model
         if new_key in model.state_dict():
@@ -395,15 +399,15 @@ def load_checkpoint(model, checkpoint_path, strict=True):
     
     # Add missing TPS parameters if needed
     tps_params = ['gridGen.P_X_base', 'gridGen.P_Y_base', 'gridGen.Li', 
-                 'gridGen.grid_X', 'gridGen.grid_Y']
+                 'gridGen.grid_X_base', 'gridGen.grid_Y_base']
     for param in tps_params:
         if param not in new_state_dict and hasattr(model, 'gridGen'):
-            print(f"Initializing missing TPS parameter: {param}")
-            # Initialize with current model's value
-            new_state_dict[param] = model.state_dict()[param]
+            if param in model.state_dict():
+                print(f"Initializing missing TPS parameter: {param}")
+                new_state_dict[param] = model.state_dict()[param]
     
     # Load the state dict
-    model.load_state_dict(new_state_dict, strict=False)  # Use strict=False to ignore missing keys
+    model.load_state_dict(new_state_dict, strict=False)
     
     # Print warnings
     model_keys = set(model.state_dict().keys())