Program crashed in the last step in test Tensorflow-gpu 2.0.0

When using Tensorflow 2.0.0 and split dataset into train-set and test-set. The training and testing code is as following:

for epoch in range(params.num_epochs):
    for step, (x_batch_train, y_batch_train) in enumerate(train_dist_dataset):
        DO TRAINING HERE....
    if epoch % params.num_epoch_record == 0:
        for step, (x_test, y_test) in enumerate(test_dist_dataset):
            DO TESTing HERE....
        checkpoint.step.assign_add(1)
        save_path = manager.save()
        logger.info("Saved checkpoint {}".format(save_path))

However, when after the last test data in enumerate(test_dist_dataset) the program will crash and shows up:

F .\tensorflow/core/kernels/conv_2d_gpu.h:964] Non-OK-status: GpuLaunchKernel( SwapDimension1And2InTensor3UsingTiles<T, kNumThreads, kTileSize, kTileSize, conjugate>, total_tiles_count, kNumThreads, 0, d.stream(), input, input_dims, output) status: Internal: invalid configuration argument

So, how it occurs and how to solve it?