在模型中使用自定义图层时，Keras load_model会导致'TypeError: Keyword参数未被理解：‘

from tensorflow.keras.layers import Layer

class CustomLayer(Layer):
    def __init__(self, units=32):
        super(CustomLayer, self).__init__()
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(shape=(input_shape[-1], self.units),
                                 initializer='random_normal',
                                 trainable=True)

    def call(self, inputs):
        return tf.matmul(inputs, self.w)

    def get_config(self):
        config = super(CustomLayer, self).get_config()
        config.update({'units': self.units})
        return config

    @classmethod
    def from_config(cls, config):
        return cls(**config)