How to vectorize the sum? tensor[i,:,:,:] + tensor[i]
I want to vectorize the following code:

def style_noise(self, y, style):
    n = torch.randn(y.shape)
    for i in range(n.shape[0]):
        n[i] = (n[i] - n.mean(dim=(1, 2, 3))[i]) * style.std(dim=(1, 2, 3))[i] / n.std(dim=(1, 2, 3))[i] + style.mean(dim=(1, 2, 3))[i]
    noise = Variable(n, requires_grad=False).to(y.device)
    return noise


I didn't find a way nice way of doing so.

y and style are 4d tensors, say  style.shape = y.shape = [64, 3, 128, 128].

I want to return the noise tensor, noise.shape = [64, 3, 128, 128].

Please let me know in the comments if the question is not clear.

Your use case is exactly why the .mean and .std methods come with a keepdim parameter. You can make use of this to enable broadcasting semantics to vectorize things for you:

def style_noise(self, y, style):
    n = torch.randn(y.shape)
    n_mean = n.mean(dim=(1, 2, 3), keepdim=True)
    n_std = n.std(dim=(1, 2, 3), keepdim=True)
    style_mean = style.mean(dim=(1, 2, 3), keepdim=True)
    style_std = style.std(dim=(1, 2, 3), keepdim=True)
    n = (n - n_mean) * style_std / n_std + style_mean
    noise = Variable(n, requires_grad=False).to(y.device)
    return noise

