如何创建一个仅在评估阶段生效的keras层(并且在训练阶段是透明的)？

from tensorflow.keras import backend as K
from tensorflow.keras.layers import Layer

class EvaluationOnlyLayer(Layer):
    def __init__(self, **kwargs):
        super(EvaluationOnlyLayer, self).__init__(**kwargs)
        self.supports_masking = True

    def call(self, inputs, training=None):
        if training:
            return inputs  # 在训练阶段直接返回输入
        else:
            return K.square(inputs)  # 在评估阶段对输入进行平方操作

    def get_config(self):
        config = super(EvaluationOnlyLayer, self).get_config()
        return config

from tensorflow.keras.models import Sequential

model = Sequential()
model.add(EvaluationOnlyLayer(input_shape=(10,)))