Code speed optimization

Record a section of the same function , Small pieces of code with vastly different speeds ：

I wrote it , Running speed ： The bearer of the tortoise world ！！！

zero = torch.zeros_like (d)
d_nonoise = torch.zeros_like (d)
for i in range (zero.shape[0]):
    for j in range (zero.shape[1]):
        d_nonoise[i][j] = torch.where (d[i][j] < (d[i][j].mean () + 3 * d[i][j].std ()), d[i][j], zero[i][j])
        d_nonoise[i][j] = torch.where (d[i][j] > (d[i][j].mean () - 3 * d[i][j].std ()), d[i][j], zero[i][j])

return d

Code after consultation ： The shoulder in the rabbit ！

d_mean = d.view (d.shape[0], d.shape[1], -1).mean (dim=-1, keepdim=True).unsqueeze (dim=-1).repeat (1, 1,d.shape[2],d.shape[3])
d_std = d.view (d.shape[0], d.shape[1], -1).std (dim=-1, keepdim=True).unsqueeze (dim=-1).repeat (1, 1,d.shape[2],d.shape[3])
zero = torch.zeros (size=(1,), device=d.device, dtype=d.dtype)
d_nonoise = torch.where ((d < (d_mean + 3 * d_std)) & (d > (d_mean - 3 * d_std)), d, zero)
d = d_nonoise
return d

Just take this as a warning ： Write good code , Write the code ！ Stop making garbage ！