在TPU上训练时，如何在tf keras中保存模型权重？

import tensorflow as tf
import os

# 检查TPU是否可用
try:
  tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
  print('Running on TPU:', tpu.master())
except ValueError:
  tpu = None

# 如果可用，则将TPU初始化为主要设备
if tpu:
  tf.config.experimental_connect_to_cluster(tpu)
  tf.tpu.experimental.initialize_tpu_system(tpu)
  strategy = tf.distribute.experimental.TPUStrategy(tpu)
else:
  strategy = tf.distribute.get_strategy()

# 在TPU上创建模型
with strategy.scope():
  model = tf.keras.Sequential([...])  # 模型的网络结构

# 编译模型
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# 加载训练数据
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size)

# 训练模型
model.fit(train_dataset, epochs=num_epochs, steps_per_epoch=steps_per_epoch)

# 创建保存模型权重的目录
os.makedirs('tpu_saved_model', exist_ok=True)

# 保存模型权重
model.save_weights('tpu_saved_model/model_weights.h5')

# 创建相同结构的模型
model = tf.keras.Sequential([...])  # 模型的网络结构

# 加载保存的模型权重
model.load_weights('tpu_saved_model/model_weights.h5')