from tensorflow.python.training import moving_averages

                if is_train:
                    # For training, do batch norm with batch mean & variance
                    # Update moving averages if training
                    (X, mean, variance) = tf.nn.fused_batch_norm(X, scale, offset, epsilon=epsilon, is_training=True)
                    update_mean = moving_averages.assign_moving_average(moving_mean, mean, decay)
                    update_variance = moving_averages.assign_moving_average(moving_variance, variance, decay)
                    with tf.control_dependencies([update_mean, update_variance]):
                        X = tf.identity(X)
                else:
                    # For prediction, do batch norm with computed moving mean & variance from training
