I’m not a GAN expert but hopefully we can figure this out 
In PyTorch, all normalization layers have the same formula (subtract mean, divide by STD). The only difference is in what you are normalizing against. See more here: In-layer normalization techniques for training very deep neural networks | AI Summer 5
The extra twist here is that you are trying to learn the normalization STD as a network parameter, which is something I have never seen before. You shouldn’t need a gradient for it! The STD is normally calculated by keeping a buffer, remembering the values you saw and computing it from there (like a running mean…). Exposing it is a parameter means that you want to compute a gradient of the loss wrt to the STD which is odd to me - but I’m not a GAN expert so I don’t know if this is something that somehow people do.
According to this part:
std = sigmas[st]
loss.append(((x[:, st] - torch.tanh(recon_x[:, st])) ** 2 / 2 / (std ** 2)).sum())
loss.append(torch.log(std) * x.size()[0])

it seems to me you are normalizing against the last dimension, so this seems akin to LayerNorm? I’m not sure I read that right though, so I’d recommend investigating deeper rather than just trusting me 
Another thing to say is that in my experience, deep nets care that you do normalize, but what layer you choose tends to have little to no impact on the eventual results. So while you dig deeper into the formula to figure out the right layer to use, I do recommend you just put a LayerNorm there and let your computer crunch while you are looking at the paper. Don’t use Opacus for this! Just replace that part and keep everything the same, run it and see if you get similar performance.