class_docstr: Megatron Vision Transformer Model.
class_docstr_tok: ["Megatron","Vision","Transformer","Model","."]
code: def __init__(self, cfg: DictConfig, trainer: Trainer):
      if not HAVE_APEX:
          raise ImportError(
              "Apex was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
          )
      if not HAVE_MEGATRON_CORE:
          raise ImportError(
              "megatron-core was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
          )
  
      super().__init__(cfg, trainer=trainer)
  
      self._validate_trainer()
  
      # TODO(yuya): clean up all default values
      self.megatron_amp_O2 = cfg.get("megatron_amp_O2", False)
  
      if not self.megatron_amp_O2 and self.cfg.get(
          "virtual_pipeline_model_parallel_size", None
      ):
          raise ValueError(
              "Virtual pipeline model parallel is only supported when using megatron_amp_O2"
          )
  
      # build_model returns a list of modules which are used for interleaved pipeline parallelism
      if isinstance(self.trainer.accelerator, CPUAccelerator):
          self.model = build_model(
              model_provider_func=self.model_provider_func,
              wrap_with_ddp=False,
              on_cpu=True,
              virtual_pipeline_model_parallel_size=self.cfg.get(
                  "virtual_pipeline_model_parallel_size", None
              ),
          )
      else:
          self.model = build_model(
              model_provider_func=self.model_provider_func,
              wrap_with_ddp=False,
              virtual_pipeline_model_parallel_size=self.cfg.get(
                  "virtual_pipeline_model_parallel_size", None
              ),
          )
  
      # if we're not using interleaved, then self.model is a module.
      if (
          self.cfg.get("virtual_pipeline_model_parallel_size", None)
          is None
      ):
          self.model = self.model[0]
  
      if self.megatron_amp_O2:
          if not self.with_distributed_adam:
              # Pre-allocate the model on GPU to have master parameters allocated on the same device with matching data type
              if isinstance(self.model, list):
                  for module in self.model:
                      module.cuda(torch.cuda.current_device())
              else:
                  self.model.cuda(torch.cuda.current_device())
  
          # Model wrapper to convert both model and inputs to half precision
          if isinstance(self.model, list):
              converted_model = []
              for module in self.model:
                  converted_model.append(
                      Float16Module(
                          config=self.model_parallel_config,
                          module=module,
                          precision=cfg.precision,
                      )
                  )
                  self.model = converted_model
          else:
              self.model = Float16Module(
                  config=self.model_parallel_config,
                  module=self.model,
                  precision=cfg.precision,
              )
      self.autocast_dtype = torch_dtype_from_precision(
          self.trainer.precision
      )
      self.enable_autocast = (
          True
          if (not self.megatron_amp_O2)
          and (self.autocast_dtype in [torch.float16, torch.bfloat16])
          else False
      )
  
      self.transformer_engine = cfg.get("transformer_engine", False)
  
      # Convert the global-batch-based profile index to micro-batch index
      if hasattr(self, "_nsys_profile_enabled"):
          mp_size = cfg.get("tensor_model_parallel_size", 1) * cfg.get(
              "pipeline_model_parallel_size", 1
          )
          data_parallel_world_size = trainer.world_size // mp_size
          grad_accum_steps = cfg.get("global_batch_size") // (
              cfg.get("micro_batch_size") * data_parallel_world_size
          )
          self._nsys_profile_start_step *= grad_accum_steps
          self._nsys_profile_end_step *= grad_accum_steps
      self.get_attention_mask_from_fusion = self.cfg.get(
          "get_attention_mask_from_fusion", True
      )
      self.initialize_ub = self.cfg.get("ub_tp_comm_overlap", False)
  
imports: import itertools
  from functools import partial
  from typing import Any, Optional
  import torch
  from omegaconf.dictconfig import DictConfig
  from pytorch_lightning.accelerators import CPUAccelerator
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.nlp.data.language_modeling.megatron.data_samplers import MegatronPretrainingSampler
  from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel
  from nemo.collections.nlp.modules.common.megatron.build_model import build_model
  from nemo.collections.nlp.modules.common.megatron.module import Float16Module, MegatronModule
  from nemo.collections.nlp.modules.common.megatron.utils import average_losses_across_data_parallel_group, get_all_params_for_weight_decay_optimization, get_linear_layer, get_params_for_weight_decay_optimization, init_method_normal, scaled_init_method_normal
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank, torch_dtype_from_precision
  from nemo.collections.vision.data.megatron.data_samplers import MegatronVisionPretrainingRandomSampler
  from nemo.collections.vision.data.megatron.vit_dataset import build_train_valid_datasets
  from nemo.collections.vision.modules.vit.vit_backbone import VitBackbone, VitMlpHead
  from nemo.core.classes.common import PretrainedModelInfo
  from nemo.utils import logging
label_desc: Use image node_id 1 to create a new MegatronVitClassificationModel object from inherited base classes: MegatronBaseModel with example: obj = MegatronVitClassificationModel(cfg, trainer)