如何在TF2中构建带公差的自定义精度指标？

import tensorflow as tf
from tensorflow.keras.metrics import Metric

class CustomAccuracy(Metric):
    def __init__(self, name='custom_accuracy', **kwargs):
        super(CustomAccuracy, self).__init__(name=name, **kwargs)
        self.total = self.add_weight('total', initializer='zeros')
        self.count = self.add_weight('count', initializer='zeros')

    def update_state(self, y_true, y_pred, sample_weight=None):
        # 计算精度
        accuracy = tf.reduce_mean(tf.cast(tf.equal(y_true, y_pred), tf.float32))
        # 更新总精度和样本数量
        self.total.assign_add(accuracy)
        self.count.assign_add(1.0)

    def result(self):
        # 计算平均精度
        return self.total / self.count

class CustomLoss(tf.keras.losses.Loss):
    def __init__(self, name='custom_loss', **kwargs):
        super(CustomLoss, self).__init__(name=name, **kwargs)

    def call(self, y_true, y_pred):
        # 计算损失
        loss = tf.reduce_mean(tf.square(y_true - y_pred))
        return loss

model = tf.keras.Sequential([...])  # 构建模型
model.compile(optimizer='adam', loss=CustomLoss(), metrics=[CustomAccuracy()])  # 编译模型

model.fit(x_train, y_train, epochs=10, batch_size=32)