How to register gradient with TF2

register_gradient.py

@tf.RegisterGradient("GuidedRelu")
def _GuidedReluGrad(op, grad):
    gate_f = tf.cast(op.outputs[0] > 0, "float32")  # Filter must be activated
    gate_R = tf.cast(grad > 0, "float32")  # Grads must be positive
    return gate_f * gate_R * grad

with tf.Graph().as_default() as g:
    model = tf.keras.applications.resnet50.ResNet50(weights='imagenet', include_top=True)

    with g.gradient_override_map({"Relu": "GuidedRelu"}):
        # Do stuff here