在Tensorflow联合成员中使用“accuracy”时获取“TypeError: Caught”

import tensorflow as tf
import tensorflow_federated as tff

# 定义模型
def create_model():
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(10, activation='relu', input_shape=(784,)),
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    return model

# 定义训练过程
@tff.tf_computation
def train_on_batch(model, batch):
    with tf.GradientTape() as tape:
        logits = model(batch['x'])
        loss = tf.reduce_mean(
            tf.nn.sparse_softmax_cross_entropy_with_logits(
                labels=batch['y'], logits=logits))
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

# 定义评估指标
@tff.tf_computation
def evaluate(model, dataset):
    predictions = model(dataset['x'])
    accuracy = tf.reduce_mean(
        tf.keras.metrics.sparse_categorical_accuracy(
            y_true=dataset['y'], y_pred=predictions))
    return accuracy

# 在本地设备上计算评估指标
local_evaluate = tff.federated_computation(evaluate, model)

# 在服务器上聚合评估指标
federated_evaluate = tff.federated_mean(local_evaluate)