docstr: This function should only be called on lightning. It handles checkpoint
  update that is being called by LightningModule's `on_load_checkpoint`,
  which should update the checkpoint to the format desired. The logic
  contains two parts, when checkpoint is a model only checkpoint and
  when checkpoint is a trainer checkpoint. This function applies the checkpoint
  update in place.
  
  If the checkpoint is a model only checkpoint:
      1. If it is an mmf checkpoint, convert to lightning format
          putting it inside a "state_dict" key
      2. Apply the model's format state key to give the model a chance to update
      3. If config.checkpoint.pretrained_state_mapping is True, apply
          the mapping speicified in the config, and remove the keys that exist
          in the checkpoint that do not exist in the mapping.
  The updated checkpoint should be of the format: {"state_dict": ckpts}, where
  ckpts should be the model state_dict.
  
  If the checkpoint is a trainer only checkpoint:
      1. do the above steps for model checkpoint update
      2. do the checkpoint trainer state update from mmf to lightning
  The updated checkpoint should be of the format: {
      `epoch`: x,
      `global_step`: x,
      `pytorch-lightning_version`: x,
      `state_dict`: x,
      `callbacks`: x,
      `optimizer_states`: [x],
      `lr_schedulers`: [x],
  }
docstr_tok: ["This","function","should","only","be","called","on","lightning",".","It","handles","checkpoint","update","that","is","being","called","by","LightningModule","'s","`","on_load_checkpoint","`",",","which","should","update","the","checkpoint","to","the","format","desired",".","The","logic","contains","two","parts",",","when","checkpoint","is","a","model","only","checkpoint","and","when","checkpoint","is","a","trainer","checkpoint",".","This","function","applies","the","checkpoint","update","in","place",".","If","the","checkpoint","is","a","model","only","checkpoint",":","1",".","If","it","is","an","mmf","checkpoint",",","convert","to","lightning","format","putting","it","inside","a","``","state_dict","''","key","2",".","Apply","the","model","'s","format","state","key","to","give","the","model","a","chance","to","update","3",".","If","config.checkpoint.pretrained_state_mapping","is","True",",","apply","the","mapping","speicified","in","the","config",",","and","remove","the","keys","that","exist","in","the","checkpoint","that","do","not","exist","in","the","mapping",".","The","updated","checkpoint","should","be","of","the","format",":","{","``","state_dict","''",":","ckpts","}",",","where","ckpts","should","be","the","model","state_dict",".","If","the","checkpoint","is","a","trainer","only","checkpoint",":","1.","do","the","above","steps","for","model","checkpoint","update","2.","do","the","checkpoint","trainer","state","update","from","mmf","to","lightning","The","updated","checkpoint","should","be","of","the","format",":","{","`","epoch","`",":","x",",","`","global_step","`",":","x",",","`","pytorch-lightning_version","`",":","x",",","`","state_dict","`",":","x",",","`","callbacks","`",":","x",",","`","optimizer_states","`",":","[","x","]",",","`","lr_schedulers","`",":","[","x","]",",","}"]
code: def update_checkpoint(
      self, checkpoint: Dict[str, Any], model: torch.nn.Module
  ) -> None:
      r"""
      This function should only be called on lightning. It handles checkpoint
      update that is being called by LightningModule's `on_load_checkpoint`,
      which should update the checkpoint to the format desired. The logic
      contains two parts, when checkpoint is a model only checkpoint and
      when checkpoint is a trainer checkpoint. This function applies the checkpoint
      update in place.
  
      If the checkpoint is a model only checkpoint:
          1. If it is an mmf checkpoint, convert to lightning format
              putting it inside a "state_dict" key
          2. Apply the model's format state key to give the model a chance to update
          3. If config.checkpoint.pretrained_state_mapping is True, apply
              the mapping speicified in the config, and remove the keys that exist
              in the checkpoint that do not exist in the mapping.
      The updated checkpoint should be of the format: {"state_dict": ckpts}, where
      ckpts should be the model state_dict.
  
      If the checkpoint is a trainer only checkpoint:
          1. do the above steps for model checkpoint update
          2. do the checkpoint trainer state update from mmf to lightning
      The updated checkpoint should be of the format: {
          `epoch`: x,
          `global_step`: x,
          `pytorch-lightning_version`: x,
          `state_dict`: x,
          `callbacks`: x,
          `optimizer_states`: [x],
          `lr_schedulers`: [x],
      }
      """
  
      if is_model_only_checkpoint(checkpoint):
          self._update_model_checkpoint(
              checkpoint=checkpoint, model=model
          )
          return
  
      # this assumes the checkpoint is trainer only
      if not is_pl_trainer_checkpoint(checkpoint):
          self._update_trainer_checkpoint_from_mmf(
              checkpoint=checkpoint, model=model
          )
  
file_path: mmf/mmf/utils/checkpoint_updater.py
imports: import logging
  from typing import Any, Dict, Tuple
  import torch
  from mmf.common.registry import registry
label_desc: Use image node_id 2 for calling the MMFToPLCheckpointUpdater obj's underlying member method code with example usage: obj.update_checkpoint(checkpoint, model) without return types