I am doing something similar with an exponential distribution. For some reason every element in my batch has the same scale and shift output by the encoder.  Am I getting i.i.d. samples across each dimension if I sample like this? Note z.shape = [ kl_sample_count, batch_size, feature_size]:
prior = torch.distributions.normal.Normal(torch.zeros((batch_size, feature_size)), torch.ones((batch_size, feature_size)))
z = prior.rsample(torch.Size([kl_sample_count]))

EDIT: I have singled out the cause as the hidden layer equaling all zeros after the ReLU, causing every output to equal the bias of the final layer. But still not sure what’s causing this…