db9d883aecb6cdfba6c6bbc76b83d85397fef28d,maml_rl/utils/torch_utils.py,,weighted_mean,#Any#Any#Any#,4

Before Change

    if isinstance(dim, int):
        dim = (dim,)
    numerator = tensor * weights
    denominator = weights
    for dimension in dim:
        numerator = torch.sum(numerator, dimension, keepdim=True)
        denominator = torch.sum(denominator, dimension, keepdim=True)
    return numerator / denominator

def detach_distribution(pi):
    if isinstance(pi, Categorical):
        distribution = Categorical(logits=pi.logits.detach())

After Change

        out = torch.mean(tensor)
    if dim is None:
        out = torch.sum(tensor * weights)
        out.div_(torch.sum(weights))
    else:
        mean_dim = torch.sum(tensor * weights, dim=dim)
        mean_dim.div_(torch.sum(weights, dim=dim))
        out = torch.mean(mean_dim)

Instances