import tensorflow as tf

  weighted_loss = loss_weight * loss
  summaries = {
      'regularization_loss/KL/PriorMean/Mean':
          tf.math.reduce_mean(tf.constant(prior_mean)),
      'regularization_loss/KL/PriorVar/Mean':
          tf.math.reduce_mean(tf.constant(prior_stddev)**2),
