code: def compute_lambda_ndcg(
      prediction: torch.Tensor,
      label: torch.Tensor,
      weight: torch.Tensor,
      session_lengths: torch.Tensor,
      use_exp_gain: bool,
  ) -> torch.Tensor:
      """
      Compute the sum lambda NDCG loss from a group of sessions.
  
      Args:
          prediction(torch.Tensor): a tensor of predicted scores
          label(torch.Tensor): a tensor of labels
          weight(torch.Tensor): a tensor of weights
          session_lengths(torch.Tensor): a tensor of session lengths converted from
              session_ids
          use_exp_gain(bool): whether to use exponential gain or not
  
      Returns:
          sum_loss(torch.Tensor): a tensor of the sum of the ndcg loss
      """
  
      loss = torch.zeros_like(session_lengths, dtype=torch.double)
  
      cur_index = int(0)
      for i, session_length in enumerate(session_lengths):
          data_indexes = torch.arange(
              cur_index,
              cur_index + int(session_length),
              dtype=torch.long,
              device=prediction.device,
          )
          session_loss = compute_lambda_ndcg_by_session(
              prediction=torch.take(prediction, data_indexes),
              label=torch.take(label, data_indexes),
              use_exp_gain=use_exp_gain,
          )
          loss[i] = session_loss * torch.max(
              torch.take(weight, data_indexes)
          )
          cur_index += session_length
  
      return torch.sum(loss)
  
imports: from typing import Any, cast, Dict, List, Optional, Type
  import torch
  from torch import distributed
  from torchrec.metrics.metrics_config import RecComputeMode, RecTaskInfo
  from torchrec.metrics.metrics_namespace import MetricName, MetricNamespace, MetricPrefix
  from torchrec.metrics.rec_metric import MetricComputationReport, RecMetric, RecMetricComputation, RecMetricException
label_desc: Use image node_id 2 for calling a global function with example usage: compute_lambda_ndcg(prediction, label, weight, session_lengths, use_exp_gain) and returns: torch