code: def loss_func(
      images, masks, masked_images, outputs, non_loss_data=False
  ):
      outputs = outputs.contiguous().float()
      masks_flip = 1 - masks
      flip_masked_outputs = outputs.masked_fill(masks_flip.bool(), 0)
      flip_masked_images = images.masked_fill(masks_flip.bool(), 0)
  
      ssim_fun = SSIM()
      psnr_fun = PSNR()
  
      if not non_loss_data:
          mask_count = torch.count_nonzero(masks)
          loss = F.mse_loss(
              flip_masked_outputs,
              flip_masked_images.float(),
              reduction="sum",
          )
          loss = loss / mask_count
          ssim = ssim_fun(
              flip_masked_outputs, flip_masked_images.float()
          )
          psnr = psnr_fun(
              flip_masked_outputs, flip_masked_images.float()
          )
  
          averaged_loss = average_losses_across_data_parallel_group(
              [loss, psnr, ssim]
          )
  
          return loss, {
              "loss": averaged_loss[0],
              "psnr": averaged_loss[1],
              "ssim": averaged_loss[2],
          }
      else:
          synth_images = masked_images.float() + flip_masked_outputs
          ssim = ssim_fun(synth_images, images.float())
          psnr = psnr_fun(synth_images, images.float())
          return (
              torch.cat((images, masked_images, synth_images), dim=2),
              ssim,
              psnr,
          )
  
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 3 for calling a global function with example usage: loss_func(images, masks, masked_images, outputs, non_loss_data) and returns: loss,  dict, torch,  ssim,  psnr