class_docstr: Megatron Vision Transformer Model.
class_docstr_tok: ["Megatron","Vision","Transformer","Model","."]
code: def on_validation_epoch_end(self):
      # TODO (yuya): need fix later, check with Sean
      if not self.validation_step_outputs:
          return None
  
      if parallel_state.is_pipeline_last_stage():
          loss_outputs = [
              output[0] for output in self.validation_step_outputs
          ]
          acc_outputs = [
              output[1] for output in self.validation_step_outputs
          ]
  
          averaged_metrics = torch.tensor(
              [
                  torch.stack(loss_outputs).mean(),
                  torch.stack(acc_outputs).mean(),
              ],
              dtype=torch.float32,
              device="cuda",
          )
      else:
          averaged_metrics = torch.tensor(
              [0.0, 0.0], dtype=torch.float32, device="cuda"
          )
  
      # we can only log on one rank if it is rank zero so we broadcast from last rank
      torch.distributed.broadcast(averaged_metrics, get_last_rank())
  
      averaged_loss, averaged_acc = averaged_metrics
  
      self.log(
          "global_step",
          self.trainer.global_step,
          prog_bar=True,
          rank_zero_only=True,
          batch_size=1,
      )
      self.log(
          "val_loss",
          averaged_loss,
          prog_bar=True,
          rank_zero_only=True,
          batch_size=1,
      )
      self.log(
          "val_accuracy",
          averaged_acc,
          prog_bar=True,
          rank_zero_only=True,
          batch_size=1,
      )
      self.validation_step_outputs.clear()  # free memory
  
      return averaged_loss
  
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 17 for calling the MegatronVitClassificationModel obj's underlying member method code with example usage: obj.on_validation_epoch_end() and returns: averaged_loss, None