bert embedding取平均

def get_mean_emb(hidden_states, mask):
    s = torch.sum(hidden_states * mask.unsqueeze(-1).float(), dim=1)
    d = mask.sum(axis=1, keepdim=True).float()
    return s / d