code: def compute_latency_loss(self, model, sample, net_output):
      assert (
          net_output[-1].encoder_padding_mask is None
          or not net_output[-1].encoder_padding_mask[:, 0].any()
      ), "Only right padding on source is supported."
      # 1. Obtain the expected alignment
      alpha_list = [item["alpha"] for item in net_output[1].attn_list]
      num_layers = len(alpha_list)
      bsz, num_heads, tgt_len, src_len = alpha_list[0].size()
  
      # bsz * num_layers * num_heads, tgt_len, src_len
      alpha_all = torch.cat(alpha_list, dim=1).view(
          -1, tgt_len, src_len
      )
  
      # 2 compute expected delays
      # bsz * num_heads * num_layers, tgt_len, src_len for MMA
      steps = (
          torch.arange(1, 1 + src_len)
          .unsqueeze(0)
          .unsqueeze(1)
          .expand_as(alpha_all)
          .type_as(alpha_all)
      )
  
      expected_delays = torch.sum(steps * alpha_all, dim=-1)
  
      target_padding_mask = (
          model.get_targets(sample, net_output)
          .eq(self.padding_idx)
          .unsqueeze(1)
          .expand(bsz, num_layers * num_heads, tgt_len)
          .contiguous()
          .view(-1, tgt_len)
      )
  
      src_lengths = (
          sample["net_input"]["src_lengths"]
          .unsqueeze(1)
          .expand(bsz, num_layers * num_heads)
          .contiguous()
          .view(-1)
      )
      expected_latency = LATENCY_METRICS[self.latency_avg_type](
          expected_delays,
          src_lengths,
          None,
          target_padding_mask=target_padding_mask,
      )
  
      # 2.1 average expected latency of heads
      # bsz, num_layers * num_heads
      expected_latency = expected_latency.view(bsz, -1)
      if self.latency_gather_method == "average":
          # bsz * tgt_len
          expected_latency = expected_delays.mean(dim=1)
      elif self.latency_gather_method == "weighted_average":
          weights = torch.nn.functional.softmax(expected_latency, dim=1)
          expected_latency = torch.sum(
              expected_latency * weights, dim=1
          )
      elif self.latency_gather_method == "max":
          expected_latency = expected_latency.max(dim=1)[0]
      else:
          raise NotImplementedError
  
      expected_latency = expected_latency.sum()
      avg_loss = self.latency_avg_weight * expected_latency
  
      # 2.2 variance of expected delays
      expected_delays_var = (
          expected_delays.view(bsz, -1, tgt_len).var(dim=1).mean(dim=1)
      )
      expected_delays_var = expected_delays_var.sum()
      var_loss = self.latency_avg_weight * expected_delays_var
  
      # 3. Final loss
      latency_loss = avg_loss + var_loss
  
      return latency_loss, expected_latency, expected_delays_var
  
file_path: fairseq/fairseq/criterions/label_smoothed_cross_entropy_latency_augmented.py
imports: from dataclasses import dataclass, field
  import torch
  from fairseq import utils
  from fairseq.logging import metrics
  from fairseq.criterions import register_criterion
  from fairseq.criterions.label_smoothed_cross_entropy import LabelSmoothedCrossEntropyCriterion, LabelSmoothedCrossEntropyCriterionConfig
label_desc: Use image node_id 3 for calling the LatencyAugmentedLabelSmoothedCrossEntropyCriterion obj's underlying member method code with example usage: obj.compute_latency_loss(model, sample, net_output) and returns: latency_loss,  expected_latency,  expected_delays_var