pytorch1.1.0和1.3.1的版本之间一个小差异(torch.mean(output, axis=1))

pytorch1.3.1:

def compactness_loss(output):
    # dim: (batch, vector)
    _, n = output.size()
    avg = torch.mean(output, axis=1)
    std = torch.std(output, axis=1)
    zt = output.T - avg
    zt /= std
    corr = torch.matmul(zt.T, zt) / (n - 1)
    loss = torch.sum(torch.triu(corr, diagonal=1)**2)
    return loss

但是pytorch1.1.0如果想使用上述代码只能改成:

def compactness_loss(output):
    # dim: (batch, vector)
    _, n = output.size()

    avg = torch.mean(output, dim=1)
    std = torch.std(output, dim=1)
    #zt = output.T - avg
    zt = output.t() - avg
    zt /= std
    #corr = torch.matmul(zt.T, zt) / (n - 1)
    corr = torch.matmul(zt.t(), zt) / (n - 1)
    loss = torch.sum(torch.triu(corr, diagonal=1)**2)
    return loss