def cal_metric(delays, src_lens, tgt_lens, target_padding_mask):
        """
        delays : bsz, num_heads_x_layers, tgt_len
        src_lens : bsz, 1
        target_lens : bsz, 1
        target_padding_mask: bsz, tgt_len or None
        """
        if delays.size(1) == 1:
            return delays.new_zeros([1])

        variance_delays = delays.var(dim=1)

        if target_padding_mask is not None:
            variance_delays.masked_fill_(target_padding_mask, 0)

        return variance_delays.sum(dim=1, keepdim=True) / tgt_lens

        delays : bsz, num_heads_x_layers, tgt_len
        src_lens : bsz, 1
        target_lens : bsz, 1
        target_padding_mask: bsz, tgt_len or None
        
decoding/IAD/fairseq/examples/simultaneous_translation/utils/latency.py