如何在Tensorflow数据集中可视化Nan

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

dataset = tf.keras.datasets.mnist.load_data()
(x_train, y_train), (x_test, y_test) = dataset

# 检查NaN值
nan_indices = np.argwhere(np.isnan(x_train))

# 可视化NaN值
if len(nan_indices) > 0:
    plt.figure(figsize=(10, 5))
    for i, index in enumerate(nan_indices):
        plt.subplot(2, 5, i+1)
        plt.imshow(x_train[index[0]], cmap='gray')
        plt.title(f'Index: {index[0]}')
        plt.axis('off')
    plt.tight_layout()
    plt.show()
else:
    print("No NaN values found in the dataset.")