Layer normalization uses all the activations per instance from the batch for normalization and batch normalization uses the whole batch for each activations. Ok, but you didn’t normalize per neuron, so it was a mix of both. So we were both right and wrong.   (sorry for the confusion)
When I didn’t miss something you should use
def forward(self, z):
        mu = torch.mean(z, dim=1)
        sigma = torch.std(z, dim=1)

if your input is (B,N) B batch size and N input size, if you want to do layer normalization.