Val loss behaves strange while using custom training loop in tensorflow 2.0

I'm using a VGG16 model written in tf2.0 to train on my own datasets. Some BatchNormalization layers were included in the model and the "training" argument were set to True during training time and False during validation time as described in many tutorials. The train_loss decreased to a certain level during training as expected.
However, the val_loss behaves really strange. I checked out the output of the model after training and found out that, if I set the training argument to True, the output is quite correct, but if I set it to False, the result is incorrect at all. According to the tutorials in tensorflow website, when training is set to False , the model will normalize its inputs using the mean and variance of its moving statistics learned during training but it doesn't seem so. Am I missing something?

I've provided the trainning and validation code in the below.

def train():
    logging.basicConfig(level=logging.INFO)
    tdataset = tf.data.Dataset.from_tensor_slices((train_img_list[:200], train_label_list[:200]))
    tdataset = tdataset.map(parse_function, 3).shuffle(buffer_size=200).batch(batch_size).repeat(repeat_times)
    vdataset = tf.data.Dataset.from_tensor_slices((val_img_list[:100], val_label_list[:100]))
    vdataset = vdataset.map(parse_function, 3).batch(batch_size)

    ### Vgg model
    model = VGG_PR(num_classes=num_label)

    logging.info('Model loaded')

    start_epoch = 0
    latest_ckpt = tf.train.latest_checkpoint(os.path.dirname(ckpt_path))
    if latest_ckpt:
        start_epoch = int(latest_ckpt.split('-')[1].split('.')[0])
        model.load_weights(latest_ckpt)
        logging.info('model resumed from: {}, start at epoch: {}'.format(latest_ckpt, start_epoch))
    else:
        logging.info('training from scratch since weights no there')

    ######## training loop ########
    loss_object = tf.keras.losses.MeanSquaredError()
    val_loss_object = tf.keras.losses.MeanSquaredError()
    optimizer = tf.keras.optimizers.Adam(learning_rate=initial_lr)
    train_loss = tf.metrics.Mean(name='train_loss') 
    val_loss = tf.metrics.Mean(name='val_loss')
    writer = tf.summary.create_file_writer(log_path.format(case_num))

    with writer.as_default():
        for epoch in range(start_epoch, total_epoch):
            print('start training')
            try:
                for batch, data in enumerate(tdataset):
                    images, labels = data
                    with tf.GradientTape() as tape:
                        pred = model(images, training=True)
                        if len(pred.shape) == 2:
                            pred = tf.reshape(pred,[-1, 1, 1, num_label])
                        loss = loss_object(pred, labels)
                    gradients = tape.gradient(loss, model.trainable_variables)
                    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
                    if batch % 20 ==0:
                        logging.info('Epoch: {}, iter: {}, loss:{}'.format(epoch, batch, loss.numpy()))
                    tf.summary.scalar('train_loss', loss.numpy(), step=epoch*1250*repeat_times+batch)      # the tdataset has been repeated 5 times..
                    tf.summary.text('Zernike_coe_pred', tf.as_string(tf.squeeze(pred)), step=epoch*1250*repeat_times+batch)
                    tf.summary.text('Zernike_coe_gt', tf.as_string(tf.squeeze(labels)), step=epoch*1250*repeat_times+batch)

                    writer.flush()
                    train_loss(loss)
                model.save_weights(ckpt_path.format(epoch=epoch))
            except KeyboardInterrupt:
                logging.info('interrupted.')
                model.save_weights(ckpt_path.format(epoch=epoch))
                logging.info('model saved into {}'.format(ckpt_path.format(epoch=epoch)))
                exit(0)
            # validation step
            for batch, data in enumerate(vdataset):
                images, labels = data
                val_pred = model(images, training=False)
                if len(val_pred.shape) == 2:
                    val_pred = tf.reshape(val_pred,[-1, 1, 1, num_label])
                v_loss = val_loss_object(val_pred, labels)
                val_loss(v_loss)
            logging.info('Epoch: {}, average train_loss:{}, val_loss: {}'.format(epoch, train_loss.result(), val_loss.result()))
            tf.summary.scalar('val_loss', val_loss.result(), step = epoch)
            writer.flush()
            train_loss.reset_states()
            val_loss.reset_states()
        model.save_weights(ckpt_path.format(epoch=epoch))

The train losss reduced to a very small value like the groundtruth label are in the range of [0, 1] and the average train loss can be 0.007, but the val loss is much higher than this. The output of the model tends to be close to 0 if I set training to False.

updated on Nov. 6th: I have found an interesting thing that if I use tf.function to decorate my model in its call method, the val loss will turn to be correct, but I'm not sure what has happened?