def METHOD_NAME(attention_mask_list):
    def attn_mask_postprocess(attn_mask):
        # [b, 1, s, s]
        extended_attention_mask = attn_mask.unsqueeze(1)
        return extended_attention_mask
    return [attn_mask_postprocess(attn_mask) for attn_mask in attention_mask_list]