code: def loss_func(loss_mask: Tensor, output_tensor: Tensor):
      """Loss function.
  
      Args:
          loss_mask (Tensor): Used to mask out some portions of the loss
          output_tensor (Tensor): The tensor with the losses
      """
      args = get_args()
  
      losses = output_tensor.float()
      loss_mask = loss_mask.view(-1).float()
      if args.context_parallel_size > 1:
          loss = torch.cat(
              [
                  torch.sum(losses.view(-1) * loss_mask).view(1),
                  loss_mask.sum().view(1),
              ]
          )
          torch.distributed.all_reduce(
              loss, group=mpu.get_context_parallel_group()
          )
          loss = loss[0] / loss[1]
      else:
          loss = (
              torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()
          )
  
      # Check individual rank losses are not NaN prior to DP all-reduce.
      if args.check_for_nan_in_loss_and_grad:
          global_rank = torch.distributed.get_rank()
          assert not loss.isnan(), (
              f"Rank {global_rank}: found NaN in local forward loss calculation. "
              f"Device: {torch.cuda.current_device()}, node: {os.uname()[1]}"
          )
  
      # Reduce loss for logging.
      averaged_loss = average_losses_across_data_parallel_group([loss])
  
      return loss * args.context_parallel_size, {
          "lm loss": averaged_loss[0]
      }
  
imports: import os
  import torch
  from torch import Tensor
  from functools import partial
  from typing import Union
  from megatron import get_args
  from megatron import print_rank_0
  from megatron import get_timers
  from megatron import get_tokenizer
  from megatron.core import mpu, tensor_parallel
  from megatron.core.enums import ModelType
  from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
  from megatron.core.datasets.gpt_dataset import GPTDatasetConfig
  from megatron.core.datasets.gpt_dataset import GPTDataset
  import megatron.model
  from megatron.core.models.gpt import GPTModel
  from megatron.training import pretrain
  from megatron.core.transformer.spec_utils import import_module
  from megatron.utils import get_batch_on_this_cp_rank, get_batch_on_this_tp_rank, average_losses_across_data_parallel_group
  from megatron.arguments import core_transformer_config_from_args
  from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
label_desc: Use image node_id 3 for calling a global function with example usage: loss_func(loss_mask, output_tensor) and returns: unknown,  dict