Update networks.py

networks.py

@@ -12,7 +12,7 @@ class Options:
         # Default values
         self.fine_height = 256
         self.fine_width = 192
-        self.grid_size = 3
+        self.grid_size = 5
         self.use_dropout = False
 
 def weights_init_normal(m):
@@ -499,7 +499,30 @@ def save_checkpoint(model, save_path):
         os.makedirs(os.path.dirname(save_path))
     torch.save(model.state_dict(), save_path)
 
-def load_checkpoint(model, checkpoint_path):
+def load_checkpoint(model, checkpoint_path, strict=True):
     if not os.path.exists(checkpoint_path):
-        return
-    model.load_state_dict(torch.load(checkpoint_path, map_location=torch.device('cpu')))
+        raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
+    
+    # Load checkpoint with strict=False to ignore size mismatches
+    state_dict = torch.load(checkpoint_path, map_location=torch.device('cpu'))
+    
+    # Filter out size-mismatched keys
+    model_state_dict = model.state_dict()
+    filtered_state_dict = {k: v for k, v in state_dict.items() 
+                         if k in model_state_dict and v.size() == model_state_dict[k].size()}
+    
+    # Load the filtered state dict
+    model.load_state_dict(filtered_state_dict, strict=strict)
+    
+    # Print warnings for mismatched keys
+    missing_keys = [k for k in model_state_dict.keys() if k not in state_dict]
+    unexpected_keys = [k for k in state_dict.keys() if k not in model_state_dict]
+    size_mismatch_keys = [k for k in state_dict.keys() 
+                         if k in model_state_dict and state_dict[k].size() != model_state_dict[k].size()]
+    
+    if missing_keys:
+        print(f"Missing keys in checkpoint: {missing_keys}")
+    if unexpected_keys:
+        print(f"Unexpected keys in checkpoint: {unexpected_keys}")
+    if size_mismatch_keys:
+        print(f"Size mismatch for keys: {size_mismatch_keys}")