如何在tensorboard中可视化RNN层的直方图？

import tensorflow as tf
from tensorflow.keras.layers import SimpleRNN

model = tf.keras.Sequential()
model.add(SimpleRNN(units=64, input_shape=(10, 32)))  # 假设输入形状为(10, 32)

model.compile(optimizer='adam', loss='mse')
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir='./logs')

# 创建一个tf.summary.FileWriter对象
file_writer = tf.summary.create_file_writer('./logs')

# 定义一个自定义回调函数
class HistogramCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        with file_writer.as_default():
            # 获取RNN层的权重
            weights = self.model.layers[0].get_weights()[0]
            # 将权重数据写入直方图
            tf.summary.histogram('RNN_Weights', weights, step=epoch)
        file_writer.flush()

# 训练模型并使用自定义回调函数
model.fit(x_train, y_train, epochs=10, callbacks=[tensorboard_callback, HistogramCallback()])

tensorboard --logdir=./logs