def METHOD_NAME(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