def weighted_mean(x, sample_weight=None):
    """
    A simple torch weighted mean function
    """
    if sample_weight is None:
        return x.mean()
    else:
        assert x.shape == sample_weight.shape
        return (x * sample_weight).sum() / sample_weight.sum()