Suppose I have a tensor,
a, I wish to normalize to the mean and standard deviation along its last axis:
a.shape > torch.Size([B,C,X,Y])
I can achieve this using:
s_mean = torch.mean(s, axis=-1).unsqueeze(-1) s_std = torch.std(s, axis=-1).unsqueeze(-1) s_norm = (s-s_mean)/(s_std)
However, there are several entries along the axis
Y that have mean 0 and variance 0, ergo they are all the same value.
I would like to only normalize the values with a non-zero mean and standard deviation and set the remaining values to zero.
What functions to I require to achieve this?