prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  Add `strict_shape` option to CheckpointLoader
  **Is your feature request related to a problem? Please describe.**
  Currently, we don't support the transfer-learning case that load a checkpoint with same layer names but different shape.
  We can refer to below code:
  ```py
  model_3 = get_model_with_3_classes()
  state_dict_model_4 = torch.load("best_model_4.pt")
  
  @trainer.on(Events.STARTED, model_3, state_dict_model_4)
  def permissive_model_loader(model, state_dict):
      this_state_dict = model.state_dict()
      matched_state_dict = {
          k: v for k, v in state_dict.items()
          if k in this_state_dict and v.shape == this_state_dict[k].shape
      }
      model.load_state_dict(matched_state_dict, strict=False)
  
  trainer.run(...)
  ```
  
  Thanks.
  
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `monai/handlers/checkpoint_loader.py`
  Content:
  ```
  1 # Copyright 2020 - 2021 MONAI Consortium
  2 # Licensed under the Apache License, Version 2.0 (the "License");
  3 # you may not use this file except in compliance with the License.
  4 # You may obtain a copy of the License at
  5 #     http://www.apache.org/licenses/LICENSE-2.0
  6 # Unless required by applicable law or agreed to in writing, software
  7 # distributed under the License is distributed on an "AS IS" BASIS,
  8 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  9 # See the License for the specific language governing permissions and
  10 # limitations under the License.
  11 
  12 import logging
  13 from typing import TYPE_CHECKING, Dict, Optional
  14 
  15 import torch
  16 
  17 from monai.utils import exact_version, optional_import
  18 
  19 Events, _ = optional_import("ignite.engine", "0.4.4", exact_version, "Events")
  20 Checkpoint, _ = optional_import("ignite.handlers", "0.4.4", exact_version, "Checkpoint")
  21 if TYPE_CHECKING:
  22     from ignite.engine import Engine
  23 else:
  24     Engine, _ = optional_import("ignite.engine", "0.4.4", exact_version, "Engine")
  25 
  26 
  27 class CheckpointLoader:
  28     """
  29     CheckpointLoader acts as an Ignite handler to load checkpoint data from file.
  30     It can load variables for network, optimizer, lr_scheduler, etc.
  31     If saving checkpoint after `torch.nn.DataParallel`, need to save `model.module` instead
  32     as PyTorch recommended and then use this loader to load the model.
  33 
  34     Args:
  35         load_path: the file path of checkpoint, it should be a PyTorch `pth` file.
  36         load_dict: target objects that load checkpoint to. examples::
  37 
  38             {'network': net, 'optimizer': optimizer, 'lr_scheduler': lr_scheduler}
  39 
  40         name: identifier of logging.logger to use, if None, defaulting to ``engine.logger``.
  41         map_location: when loading the module for distributed training/evaluation,
  42             need to provide an appropriate map_location argument to prevent a process
  43             to step into others’ devices. If map_location is missing, torch.load will
  44             first load the module to CPU and then copy each parameter to where it was
  45             saved, which would result in all processes on the same machine using the
  46             same set of devices.
  47         strict: whether to strictly enforce that the keys in :attr:`state_dict` match the keys
  48             returned by this module's :meth:`~torch.nn.Module.state_dict` function. Default: ``True``
  49 
  50     """
  51 
  52     def __init__(
  53         self,
  54         load_path: str,
  55         load_dict: Dict,
  56         name: Optional[str] = None,
  57         map_location: Optional[Dict] = None,
  58         strict: bool = True,
  59     ) -> None:
  60         if load_path is None:
  61             raise AssertionError("must provide clear path to load checkpoint.")
  62         self.load_path = load_path
  63         if not (load_dict is not None and len(load_dict) > 0):
  64             raise AssertionError("must provide target objects to load.")
  65         self.logger = logging.getLogger(name)
  66         self.load_dict = load_dict
  67         self._name = name
  68         self.map_location = map_location
  69         self.strict = strict
  70 
  71     def attach(self, engine: Engine) -> None:
  72         """
  73         Args:
  74             engine: Ignite Engine, it can be a trainer, validator or evaluator.
  75         """
  76         if self._name is None:
  77             self.logger = engine.logger
  78         engine.add_event_handler(Events.STARTED, self)
  79 
  80     def __call__(self, engine: Engine) -> None:
  81         """
  82         Args:
  83             engine: Ignite Engine, it can be a trainer, validator or evaluator.
  84         """
  85         checkpoint = torch.load(self.load_path, map_location=self.map_location)
  86 
  87         # save current max epochs setting in the engine, don't overwrite it if larger than max_epochs in checkpoint
  88         prior_max_epochs = engine.state.max_epochs
  89         Checkpoint.load_objects(to_load=self.load_dict, checkpoint=checkpoint, strict=self.strict)
  90         if engine.state.epoch > prior_max_epochs:
  91             raise ValueError(
  92                 f"Epoch count ({engine.state.epoch}) in checkpoint is larger than "
  93                 f"the `engine.state.max_epochs` ({prior_max_epochs}) of engine. To further train from checkpoint, "
  94                 "construct trainer with `max_epochs` larger than checkpoint's epoch count. "
  95                 "To use checkpoint for inference, no need to load state_dict for the engine."
  96             )
  97         engine.state.max_epochs = prior_max_epochs
  98 
  99         self.logger.info(f"Restored all variables from {self.load_path}")
  100
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/monai/handlers/checkpoint_loader.py b/monai/handlers/checkpoint_loader.py
  --- a/monai/handlers/checkpoint_loader.py
  +++ b/monai/handlers/checkpoint_loader.py
  @@ -13,6 +13,7 @@
   from typing import TYPE_CHECKING, Dict, Optional
   
   import torch
  +import torch.nn as nn
   
   from monai.utils import exact_version, optional_import
   
  @@ -44,8 +45,12 @@
               first load the module to CPU and then copy each parameter to where it was
               saved, which would result in all processes on the same machine using the
               same set of devices.
  -        strict: whether to strictly enforce that the keys in :attr:`state_dict` match the keys
  -            returned by this module's :meth:`~torch.nn.Module.state_dict` function. Default: ``True``
  +        strict: whether to strictly enforce that the keys in `state_dict` match the keys
  +            returned by `torch.nn.Module.state_dict` function. default to `True`.
  +        strict_shape: whether to enforce the data shape of the matched layers in the checkpoint,
  +            `if `False`, it will skip the layers that have different data shape with checkpoint content.
  +            This can be useful advanced feature for transfer learning. users should totally
  +            understand which layers will have different shape. default to `True`.
   
       """
   
  @@ -56,6 +61,7 @@
           name: Optional[str] = None,
           map_location: Optional[Dict] = None,
           strict: bool = True,
  +        strict_shape: bool = True,
       ) -> None:
           if load_path is None:
               raise AssertionError("must provide clear path to load checkpoint.")
  @@ -67,6 +73,7 @@
           self._name = name
           self.map_location = map_location
           self.strict = strict
  +        self.strict_shape = strict_shape
   
       def attach(self, engine: Engine) -> None:
           """
  @@ -84,6 +91,20 @@
           """
           checkpoint = torch.load(self.load_path, map_location=self.map_location)
   
  +        if not self.strict_shape:
  +            k, _ = list(self.load_dict.items())[0]
  +            # single object and checkpoint is directly a state_dict
  +            if len(self.load_dict) == 1 and k not in checkpoint:
  +                checkpoint = {k: checkpoint}
  +
  +            # skip items that don't match data shape
  +            for k, obj in self.load_dict.items():
  +                if isinstance(obj, (nn.DataParallel, nn.parallel.DistributedDataParallel)):
  +                    obj = obj.module
  +                if isinstance(obj, torch.nn.Module):
  +                    d = obj.state_dict()
  +                    checkpoint[k] = {k: v for k, v in checkpoint[k].items() if k in d and v.shape == d[k].shape}
  +
           # save current max epochs setting in the engine, don't overwrite it if larger than max_epochs in checkpoint
           prior_max_epochs = engine.state.max_epochs
           Checkpoint.load_objects(to_load=self.load_dict, checkpoint=checkpoint, strict=self.strict)
  
verification_info: {"golden_diff": "diff --git a/monai/handlers/checkpoint_loader.py b/monai/handlers/checkpoint_loader.py\n--- a/monai/handlers/checkpoint_loader.py\n+++ b/monai/handlers/checkpoint_loader.py\n@@ -13,6 +13,7 @@\n from typing import TYPE_CHECKING, Dict, Optional\n \n import torch\n+import torch.nn as nn\n \n from monai.utils import exact_version, optional_import\n \n@@ -44,8 +45,12 @@\n             first load the module to CPU and then copy each parameter to where it was\n             saved, which would result in all processes on the same machine using the\n             same set of devices.\n-        strict: whether to strictly enforce that the keys in :attr:`state_dict` match the keys\n-            returned by this module's :meth:`~torch.nn.Module.state_dict` function. Default: ``True``\n+        strict: whether to strictly enforce that the keys in `state_dict` match the keys\n+            returned by `torch.nn.Module.state_dict` function. default to `True`.\n+        strict_shape: whether to enforce the data shape of the matched layers in the checkpoint,\n+            `if `False`, it will skip the layers that have different data shape with checkpoint content.\n+            This can be useful advanced feature for transfer learning. users should totally\n+            understand which layers will have different shape. default to `True`.\n \n     \"\"\"\n \n@@ -56,6 +61,7 @@\n         name: Optional[str] = None,\n         map_location: Optional[Dict] = None,\n         strict: bool = True,\n+        strict_shape: bool = True,\n     ) -> None:\n         if load_path is None:\n             raise AssertionError(\"must provide clear path to load checkpoint.\")\n@@ -67,6 +73,7 @@\n         self._name = name\n         self.map_location = map_location\n         self.strict = strict\n+        self.strict_shape = strict_shape\n \n     def attach(self, engine: Engine) -> None:\n         \"\"\"\n@@ -84,6 +91,20 @@\n         \"\"\"\n         checkpoint = torch.load(self.load_path, map_location=self.map_location)\n \n+        if not self.strict_shape:\n+            k, _ = list(self.load_dict.items())[0]\n+            # single object and checkpoint is directly a state_dict\n+            if len(self.load_dict) == 1 and k not in checkpoint:\n+                checkpoint = {k: checkpoint}\n+\n+            # skip items that don't match data shape\n+            for k, obj in self.load_dict.items():\n+                if isinstance(obj, (nn.DataParallel, nn.parallel.DistributedDataParallel)):\n+                    obj = obj.module\n+                if isinstance(obj, torch.nn.Module):\n+                    d = obj.state_dict()\n+                    checkpoint[k] = {k: v for k, v in checkpoint[k].items() if k in d and v.shape == d[k].shape}\n+\n         # save current max epochs setting in the engine, don't overwrite it if larger than max_epochs in checkpoint\n         prior_max_epochs = engine.state.max_epochs\n         Checkpoint.load_objects(to_load=self.load_dict, checkpoint=checkpoint, strict=self.strict)\n", "issue": "Add `strict_shape` option to CheckpointLoader\n**Is your feature request related to a problem? Please describe.**\r\nCurrently, we don't support the transfer-learning case that load a checkpoint with same layer names but different shape.\r\nWe can refer to below code:\r\n```py\r\nmodel_3 = get_model_with_3_classes()\r\nstate_dict_model_4 = torch.load(\"best_model_4.pt\")\r\n\r\n@trainer.on(Events.STARTED, model_3, state_dict_model_4)\r\ndef permissive_model_loader(model, state_dict):\r\n    this_state_dict = model.state_dict()\r\n    matched_state_dict = {\r\n        k: v for k, v in state_dict.items()\r\n        if k in this_state_dict and v.shape == this_state_dict[k].shape\r\n    }\r\n    model.load_state_dict(matched_state_dict, strict=False)\r\n\r\ntrainer.run(...)\r\n```\r\n\r\nThanks.\r\n\n", "before_files": [{"content": "# Copyright 2020 - 2021 MONAI Consortium\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#     http://www.apache.org/licenses/LICENSE-2.0\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nfrom typing import TYPE_CHECKING, Dict, Optional\n\nimport torch\n\nfrom monai.utils import exact_version, optional_import\n\nEvents, _ = optional_import(\"ignite.engine\", \"0.4.4\", exact_version, \"Events\")\nCheckpoint, _ = optional_import(\"ignite.handlers\", \"0.4.4\", exact_version, \"Checkpoint\")\nif TYPE_CHECKING:\n    from ignite.engine import Engine\nelse:\n    Engine, _ = optional_import(\"ignite.engine\", \"0.4.4\", exact_version, \"Engine\")\n\n\nclass CheckpointLoader:\n    \"\"\"\n    CheckpointLoader acts as an Ignite handler to load checkpoint data from file.\n    It can load variables for network, optimizer, lr_scheduler, etc.\n    If saving checkpoint after `torch.nn.DataParallel`, need to save `model.module` instead\n    as PyTorch recommended and then use this loader to load the model.\n\n    Args:\n        load_path: the file path of checkpoint, it should be a PyTorch `pth` file.\n        load_dict: target objects that load checkpoint to. examples::\n\n            {'network': net, 'optimizer': optimizer, 'lr_scheduler': lr_scheduler}\n\n        name: identifier of logging.logger to use, if None, defaulting to ``engine.logger``.\n        map_location: when loading the module for distributed training/evaluation,\n            need to provide an appropriate map_location argument to prevent a process\n            to step into others\u2019 devices. If map_location is missing, torch.load will\n            first load the module to CPU and then copy each parameter to where it was\n            saved, which would result in all processes on the same machine using the\n            same set of devices.\n        strict: whether to strictly enforce that the keys in :attr:`state_dict` match the keys\n            returned by this module's :meth:`~torch.nn.Module.state_dict` function. Default: ``True``\n\n    \"\"\"\n\n    def __init__(\n        self,\n        load_path: str,\n        load_dict: Dict,\n        name: Optional[str] = None,\n        map_location: Optional[Dict] = None,\n        strict: bool = True,\n    ) -> None:\n        if load_path is None:\n            raise AssertionError(\"must provide clear path to load checkpoint.\")\n        self.load_path = load_path\n        if not (load_dict is not None and len(load_dict) > 0):\n            raise AssertionError(\"must provide target objects to load.\")\n        self.logger = logging.getLogger(name)\n        self.load_dict = load_dict\n        self._name = name\n        self.map_location = map_location\n        self.strict = strict\n\n    def attach(self, engine: Engine) -> None:\n        \"\"\"\n        Args:\n            engine: Ignite Engine, it can be a trainer, validator or evaluator.\n        \"\"\"\n        if self._name is None:\n            self.logger = engine.logger\n        engine.add_event_handler(Events.STARTED, self)\n\n    def __call__(self, engine: Engine) -> None:\n        \"\"\"\n        Args:\n            engine: Ignite Engine, it can be a trainer, validator or evaluator.\n        \"\"\"\n        checkpoint = torch.load(self.load_path, map_location=self.map_location)\n\n        # save current max epochs setting in the engine, don't overwrite it if larger than max_epochs in checkpoint\n        prior_max_epochs = engine.state.max_epochs\n        Checkpoint.load_objects(to_load=self.load_dict, checkpoint=checkpoint, strict=self.strict)\n        if engine.state.epoch > prior_max_epochs:\n            raise ValueError(\n                f\"Epoch count ({engine.state.epoch}) in checkpoint is larger than \"\n                f\"the `engine.state.max_epochs` ({prior_max_epochs}) of engine. To further train from checkpoint, \"\n                \"construct trainer with `max_epochs` larger than checkpoint's epoch count. \"\n                \"To use checkpoint for inference, no need to load state_dict for the engine.\"\n            )\n        engine.state.max_epochs = prior_max_epochs\n\n        self.logger.info(f\"Restored all variables from {self.load_path}\")\n", "path": "monai/handlers/checkpoint_loader.py"}], "after_files": [{"content": "# Copyright 2020 - 2021 MONAI Consortium\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#     http://www.apache.org/licenses/LICENSE-2.0\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nfrom typing import TYPE_CHECKING, Dict, Optional\n\nimport torch\nimport torch.nn as nn\n\nfrom monai.utils import exact_version, optional_import\n\nEvents, _ = optional_import(\"ignite.engine\", \"0.4.4\", exact_version, \"Events\")\nCheckpoint, _ = optional_import(\"ignite.handlers\", \"0.4.4\", exact_version, \"Checkpoint\")\nif TYPE_CHECKING:\n    from ignite.engine import Engine\nelse:\n    Engine, _ = optional_import(\"ignite.engine\", \"0.4.4\", exact_version, \"Engine\")\n\n\nclass CheckpointLoader:\n    \"\"\"\n    CheckpointLoader acts as an Ignite handler to load checkpoint data from file.\n    It can load variables for network, optimizer, lr_scheduler, etc.\n    If saving checkpoint after `torch.nn.DataParallel`, need to save `model.module` instead\n    as PyTorch recommended and then use this loader to load the model.\n\n    Args:\n        load_path: the file path of checkpoint, it should be a PyTorch `pth` file.\n        load_dict: target objects that load checkpoint to. examples::\n\n            {'network': net, 'optimizer': optimizer, 'lr_scheduler': lr_scheduler}\n\n        name: identifier of logging.logger to use, if None, defaulting to ``engine.logger``.\n        map_location: when loading the module for distributed training/evaluation,\n            need to provide an appropriate map_location argument to prevent a process\n            to step into others\u2019 devices. If map_location is missing, torch.load will\n            first load the module to CPU and then copy each parameter to where it was\n            saved, which would result in all processes on the same machine using the\n            same set of devices.\n        strict: whether to strictly enforce that the keys in `state_dict` match the keys\n            returned by `torch.nn.Module.state_dict` function. default to `True`.\n        strict_shape: whether to enforce the data shape of the matched layers in the checkpoint,\n            `if `False`, it will skip the layers that have different data shape with checkpoint content.\n            This can be useful advanced feature for transfer learning. users should totally\n            understand which layers will have different shape. default to `True`.\n\n    \"\"\"\n\n    def __init__(\n        self,\n        load_path: str,\n        load_dict: Dict,\n        name: Optional[str] = None,\n        map_location: Optional[Dict] = None,\n        strict: bool = True,\n        strict_shape: bool = True,\n    ) -> None:\n        if load_path is None:\n            raise AssertionError(\"must provide clear path to load checkpoint.\")\n        self.load_path = load_path\n        if not (load_dict is not None and len(load_dict) > 0):\n            raise AssertionError(\"must provide target objects to load.\")\n        self.logger = logging.getLogger(name)\n        self.load_dict = load_dict\n        self._name = name\n        self.map_location = map_location\n        self.strict = strict\n        self.strict_shape = strict_shape\n\n    def attach(self, engine: Engine) -> None:\n        \"\"\"\n        Args:\n            engine: Ignite Engine, it can be a trainer, validator or evaluator.\n        \"\"\"\n        if self._name is None:\n            self.logger = engine.logger\n        engine.add_event_handler(Events.STARTED, self)\n\n    def __call__(self, engine: Engine) -> None:\n        \"\"\"\n        Args:\n            engine: Ignite Engine, it can be a trainer, validator or evaluator.\n        \"\"\"\n        checkpoint = torch.load(self.load_path, map_location=self.map_location)\n\n        if not self.strict_shape:\n            k, _ = list(self.load_dict.items())[0]\n            # single object and checkpoint is directly a state_dict\n            if len(self.load_dict) == 1 and k not in checkpoint:\n                checkpoint = {k: checkpoint}\n\n            # skip items that don't match data shape\n            for k, obj in self.load_dict.items():\n                if isinstance(obj, (nn.DataParallel, nn.parallel.DistributedDataParallel)):\n                    obj = obj.module\n                if isinstance(obj, torch.nn.Module):\n                    d = obj.state_dict()\n                    checkpoint[k] = {k: v for k, v in checkpoint[k].items() if k in d and v.shape == d[k].shape}\n\n        # save current max epochs setting in the engine, don't overwrite it if larger than max_epochs in checkpoint\n        prior_max_epochs = engine.state.max_epochs\n        Checkpoint.load_objects(to_load=self.load_dict, checkpoint=checkpoint, strict=self.strict)\n        if engine.state.epoch > prior_max_epochs:\n            raise ValueError(\n                f\"Epoch count ({engine.state.epoch}) in checkpoint is larger than \"\n                f\"the `engine.state.max_epochs` ({prior_max_epochs}) of engine. To further train from checkpoint, \"\n                \"construct trainer with `max_epochs` larger than checkpoint's epoch count. \"\n                \"To use checkpoint for inference, no need to load state_dict for the engine.\"\n            )\n        engine.state.max_epochs = prior_max_epochs\n\n        self.logger.info(f\"Restored all variables from {self.load_path}\")\n", "path": "monai/handlers/checkpoint_loader.py"}]}