I tried to implement something related to Layer/Group norm from scratch (without using F.batchnorm) and it would not run properly, so I dumbed it down all the way to InstanceNorm and even that does not work. I know there is a pytorch implementation and that one works, but for what I would like to do I need a custom implementation. It must be something obvious that I am missing because it is a really simple thing to do but I just can’t figure out what’s wrong…
class InstanceNorm(nn.Module):
    def __init__(self, *args, **kwargs):
        super(InstanceNorm, self).__init__()

    def forward(self, x):
        """

        :param x: has either shape (b, c, x, y) or shape (b, c, x, y, z)
        :return:
        """
        x2 = x.view(*x.shape[:2], -1) # x2.shape = (b, c, x*y(*z))

        #detach to stop gradients
        mn = x2.mean(-1).detach() # mn.shape = (b, c)
        sd = x2.std(-1).detach() # sd.shape = (b, c)
        mn = mn.view(*mn.shape[:2], *((len(x.shape) - 2) * [1])) # mn.shape = (b, c, 1, 1(, 1))
        sd = sd.view(*sd.shape[:2], *((len(x.shape) - 2) * [1])) # sd.shape = (b, c, 1, 1(, 1))

        x = (x - mn) / (sd + 1e-8)
        return x

Ideas?
Thanks a lot!