code: def on_load_checkpoint(self, checkpoint) -> None:
      # make sure inductor naming is consistent with checkpoint's
      inductor_enabled = self.cfg.get("inductor", False)
      state_dict = checkpoint["state_dict"]
      inductor_checkpoint = False
      for (
          k,
          v,
      ) in state_dict.items():
          if "_orig_mod" in k:
              inductor_checkpoint = True
              break
  
      if inductor_enabled and not inductor_checkpoint:
          # ckpt needs to be converted to inductor-format weights (add .orig_mod)
          logging.info("Add .orig_mod to all weight keys.")
          new_state_dict = {}
          for k, v in state_dict.items():
              idx = k.find("._orig_mod")
              new_key = k[:idx] + k[idx + len("._orig_mod") :]
              new_state_dict[new_key] = v
          checkpoint["state_dict"] = new_state_dict
      elif not inductor_enabled and inductor_checkpoint:
          # ckpt needs to be converted to non-inductor-format weights (remove .orig_mod)
          logging.info("Remove .orig_mod to all weight keys.")
          new_state_dict = {}
          for k, v in state_dict.items():
              new_key = k.replace("._orig_mod", "")
              new_state_dict[new_key] = v
          checkpoint["state_dict"] = new_state_dict
      super().on_load_checkpoint(checkpoint)
  
imports: import itertools
  from datetime import datetime
  from functools import partial
  from typing import Any
  import torch
  from omegaconf import DictConfig, open_dict
  from pytorch_lightning import Trainer
  from nemo.collections.multimodal.data.imagen.imagen_dataset import build_train_valid_datasets
  from nemo.collections.multimodal.models.text_to_image.imagen.precond import ContinousDDPMPrecond, EDMPrecond
  from nemo.collections.multimodal.modules.imagen.diffusionmodules.nets import EfficientUNetModel, UNetModel
  from nemo.collections.multimodal.modules.imagen.encoder.t5encoder import T5Encoder
  from nemo.collections.multimodal.modules.imagen.sampler.sampler import DDPMSampler, EDMSampler
  from nemo.collections.multimodal.parts.imagen.utils import random_dropout
  from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel
  from nemo.collections.nlp.modules.common.megatron.module import Float16Module
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank
  from nemo.core.classes.common import Serialization
  from nemo.utils import logging
label_desc: Use image node_id 22 for calling the MegatronImagen obj's underlying member method code with example usage: obj.on_load_checkpoint(checkpoint) without return types