def masked_average(tensor, mask):
    """ Performs masked average of a given tensor at time dim. """
    tensor_sum = (tensor * mask.float().unsqueeze(-1)).sum(1)
    tensor_mean = tensor_sum / mask.sum(-1).float().unsqueeze(-1)
    return tensor_mean