如何在TF 2.0中实现梯度反转层？

import tensorflow as tf
from tensorflow.keras import layers

class GradientReversalLayer(tf.keras.layers.Layer):
    def __init__(self):
        super(GradientReversalLayer, self).__init__()

    def call(self, inputs, gradient_lambda):
        return GradientReversalLayer.gradient_reverse(inputs, gradient_lambda)

    @staticmethod
    def gradient_reverse(x, gradient_lambda):
        y = tf.identity(x)
        def grad_fn(grad):
            return -grad * gradient_lambda
        tf.RegisterGradient("GradientReversal")(grad_fn)
        with tf.get_default_graph().gradient_override_map({"Identity": "GradientReversal"}):
            y = tf.identity(x)
        return y

model = tf.keras.Sequential()
model.add(layers.Dense(64, activation='relu'))
model.add(layers.Dense(32, activation='relu'))
model.add(GradientReversalLayer())
model.add(layers.Dense(10, activation='softmax'))