code: def model_provider(pre_process=True, post_process=True):
      """Build the model."""
      args = get_args()
      config = core_transformer_config_from_args(args)
      if args.vision_backbone_type == "vit":
          model = VitInpaintingModel(
              config=config,
              pre_process=pre_process,
              post_process=post_process,
          )
      elif args.vision_backbone_type == "mit":
          model = MitInpaintingModel(
              config=config,
              pre_process=pre_process,
              post_process=post_process,
          )
      else:
          raise Exception(
              "{} vision backbone is not supported.".format(
                  args.vision_backbone_type
              )
          )
      return model
  
file_path: megatron-lm/pretrain_vision_inpaint.py
imports: import torch
  import torch.nn.functional
  from functools import partial
  from megatron import get_args, get_timers, print_rank_0, print_rank_last
  from megatron.core.enums import ModelType
  from megatron.data.vit_dataset import build_train_valid_datasets
  from megatron.model.vision.inpainting import VitInpaintingModel
  from megatron.model.vision.inpainting import MitInpaintingModel
  from megatron.training import pretrain
  from megatron.utils import average_losses_across_data_parallel_group
  from tasks.vision.segmentation.metrics import SSIM, PSNR
  from megatron.arguments import core_transformer_config_from_args
label_desc: Use image node_id 1 for calling a global function with example usage: model_provider(pre_process, post_process) and returns: model