BatchNormalization Bijector Wrong Result with "prob" method

I am trying to implement a normalizing flow according to the RealNVP model for density estimation. First, I am trying to make it work on the "moons" toy dataset.

The model produces the expected result when not using the BatchNormalization bijector. However, when adding the BatchNormalization bijector to the model, the methods prob and log_prob return unexpected results.

Following is a code snippet setting up the model:

layers = 6
dimensions = 2
hidden_units = [512, 512]
bijectors = []

base_dist = tfd.Normal(loc=0.0, scale=1.0)  # specify base distribution

for i in range(layers):
    # Adding the BatchNormalization bijector corrupts the results
    bijectors.append(tfb.BatchNormalization())
    bijectors.append(RealNVP(input_shape=dimensions, n_hidden=hidden_units))
    bijectors.append(tfp.bijectors.Permute([1, 0]))

bijector = tfb.Chain(bijectors=list(reversed(bijectors))[:-1], name='chain_of_real_nvp')

flow = tfd.TransformedDistribution(
    distribution=tfd.Sample(base_dist, sample_shape=[dimensions]),
    bijector=bijector
)

When to BatchNormalization bijector is ommited both sampling and evaluating the probability return expected results:

However, when the BatchNormalization bijector is added, sampling is as expected but evaluating the probability seems wrong:

Because I am interested in density estimation the prob method is crucial. The full code can be found in the following jupyter notebook: https://github.com/mmsbrggr/normalizing-flows/blob/master/moons_training_rnvp.ipynb

I know that the BatchNormalization bijector behaves differently during training and inference. Could the problem be that the BN bijector is still in training mode? If so how can I move the flow to inference mode?