使用TPU吊舱训练keras模型？

Question

我想知道是否有人有在TPU吊舱上使用keras模型的例子

我有一个模型创建方法，它返回一个在TPU策略范围内编译的keras模型，正如许多将TPU与keras结合使用的示例所建议的那样。这适用于v3-8，但在使用更多内核（特别是v3-32）时会出现错误：

with strategy.scope():
    keras_model = create_model()
    optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08)
    keras_model.compile(loss='binary_crossentropy', optimizer=optimizer, metrics=['accuracy'])

运行model.fit时，它会失败，并出现以下错误：

Failed copying input tensor from /job:worker/replica:0/task:0/device:CPU:0 to /job:worker/replica:0/task:1/device:CPU:0 in order to run DatasetFromGraph: FetchOutputs node : not found [Op:DatasetFromGraph]

模型输入采用numpy数组的形式。是否需要tensorflow.data.Dataset