    def suppress_mem_on_mem_attention(
        self, B: int, num_heads: int, mem_size: int, attention_weight: Tensor
    ):
        """
        Arguments:
            - B: batch size
            - num_heads: number of attention heads
            - mem_size: size of memory bank
            - attention_weight: a [B*num_heads, T + 1, T + mem_size] vector

        Return:
            modified attention_weight with [B*num_heads, -1, :mem_size] = -inf
        """
        attention_weight[:, -1, :mem_size] = float("-inf")
        return attention_weight
Arguments:
- B: batch size
- num_heads: number of attention heads
- mem_size: size of memory bank
- attention_weight: a [B*num_heads, T + 1, T + mem_size] vector

Return:
modified attention_weight with [B*num_heads, -1, :mem_size] = -inf
fairseq/models/speech_to_text/modules/augmented_memory_attention.py