def fn_max(x, mask):
    """
    Args:
        x: TxBxD
        mask: BxT
    Return:
        y: BxD
    """
    mask = mask.t()[:, :, None].to(torch.bool)
    return x.masked_fill(~mask, -1e-8).max(0)[0]
Args:
x: TxBxD
mask: BxT
Return:
y: BxD
fairseq/models/wav2vec/wav2vec2_classification.py