def fn_mean(x, mask):
    """
    Args:
        x: TxBxD
        mask: BxT
    Return:
        y: BxD
    """
    if mask is not None:
        mask = mask.t()[:, :, None]
        return (x * mask).sum(0) / mask.sum(0)
    else:
        return x.sum(0) / x.shape[0]
Args:
x: TxBxD
mask: BxT
Return:
y: BxD
fairseq/models/wav2vec/wav2vec2_classification.py