why does KL loss increase during training?

Hi all,
I was trying to use BNN but I faced with an amazing point, the losses have increasing manner during the training process. would anyone give me some advice on this behavior?

logits = model(images)
kl_annealing=50
batch_size=10
labels_distribution = tfp.distributions.Categorical(logits=logits)
t = tf.Variable(0.0)
kl_regularizer = t / (kl_annealing * len(x_train) / batch_size)
neg_log_likelihood = -tf.reduce_mean(labels_distribution.log_prob(labels))
kl = sum(model.losses) / len(x_train) * tf.minimum(1.0, kl_regularizer)
elbo_loss = neg_log_likelihood + kl
predictions = tf.argmax(logits, axis=1)
accuracy, accuracy_update_op = tf.compat.v1.metrics.accuracy(labels=labels, predictions=predictions)
optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=0.001)
train_op = optimizer.minimize(elbo_loss)
update_step_op = tf.compat.v1.assign(t, t + 1)
init_op = tf.compat.v1.group(tf.compat.v1.global_variables_initializer(),tf.compat.v1.local_variables_initializer())
max_steps = 270#number of training steps to run
with tf.compat.v1.Session() as sess:
sess.run(init_op)
train_handle = sess.run(training_iterator.string_handle())
heldout_handle = sess.run(heldout_iterator.string_handle())
for step in range(max_steps):
_ = sess.run([train_op, accuracy_update_op,update_step_op],feed_dict={handle: train_handle})
if step % 10 == 0:
loss_value, accuracy_value,kl_value = sess.run([elbo_loss, accuracy,kl],feed_dict={handle: train_handle})
print("Step: {:>3d} Loss: {:.3f} Accuracy: {:.3f} KL: {:.3f}".format(step, loss_value, accuracy_value, kl_value))
Step: 0 Loss: 8.459 Accuracy: 0.300 KL: 1.046 Step: 10 Loss: 17.555 Accuracy: 0.164 KL: 11.510 Step: 20 Loss: 29.759 Accuracy: 0.186 KL: 21.973 Step: 30 Loss: 38.073 Accuracy: 0.203 KL: 32.437 Step: 40 Loss: 56.216 Accuracy: 0.190 KL: 42.900 Step: 50 Loss: 64.251 Accuracy: 0.190 KL: 53.364 Step: 60 Loss: 71.655 Accuracy: 0.193 KL: 63.827 Step: 70 Loss: 84.779 Accuracy: 0.187 KL: 74.291 Step: 80 Loss: 90.893 Accuracy: 0.189 KL: 84.754 Step: 90 Loss: 103.979 Accuracy: 0.192 KL: 95.218 Step: 100 Loss: 114.083 Accuracy: 0.194 KL: 105.681 Step: 110 Loss: 121.866 Accuracy: 0.193 KL: 116.145 Step: 120 Loss: 131.708 Accuracy: 0.193 KL: 126.608 Step: 130 Loss: 145.640 Accuracy: 0.190 KL: 137.072 Step: 140 Loss: 155.476 Accuracy: 0.194 KL: 147.535 Step: 150 Loss: 167.459 Accuracy: 0.189 KL: 157.999 Step: 160 Loss: 174.250 Accuracy: 0.188 KL: 168.462 Step: 170 Loss: 186.943 Accuracy: 0.191 KL: 178.926 Step: 180 Loss: 194.325 Accuracy: 0.192 KL: 189.389 Step: 190 Loss: 209.401 Accuracy: 0.197 KL: 199.853 Step: 200 Loss: 216.044 Accuracy: 0.201 KL: 210.316 Step: 210 Loss: 228.400 Accuracy: 0.204 KL: 220.780 Step: 220 Loss: 242.376 Accuracy: 0.202 KL: 231.243 Step: 230 Loss: 245.445 Accuracy: 0.204 KL: 241.707 Step: 240 Loss: 257.789 Accuracy: 0.206 KL: 252.171 Step: 250 Loss: 265.882 Accuracy: 0.209 KL: 262.634 Step: 260 Loss: 278.307 Accuracy: 0.210 KL: 273.098 Step: 270 Loss: 289.661 Accuracy: 0.212 KL: 283.561

Start by reproducing the supplied TensorFlow training loop and comparing the printed ELBO, KL, and accuracy values over the 270 steps. Read the loss construction around kl_regularizer, model.losses, and train_op; done means identifying the cause of the increasing loss and documenting a verified correction or explanation.