How to save&load tf.keras.Model/Layer with extended methods/properties

I want to extend tf.keras.Model to implement invertible neural networks, and save/load this model by tf.keras.Model.save() & tf.keras.models.load_model()

Specifically, I have written two classes:

• Invertible NN class
  • subclasses tf.keras.Model
  • has inverse() method
• Coupling layer class
  • subclasses tf.keras.layers.Layer
  • has inverse() method
  • has log_det_jacobian property (@property).

The code related to this looks as below (please refer to my github repo for the full code).

class SimpleRealNVP(K.Model):
    ...
    def inverse(self, inputs):
        h = inputs
        for layer in self.couplings:
            h = self.layer.inverse(h)
        return h
    ...

class AffineCoupling(L.Layer):
    ....
    @property
    def log_det_jacobian(self):
        return self._log_det_jacobian

    def inverse(self, input):
        ...

I just tried tf.keras.Model.save() and tf.keras.load_model(), but I got an error that said

...
mdl = K.models.load_model("path/to/savedmodel")
...
x_manifold = mdl.inverse(z_manifold)

# Traceback (most recent call last):
#   File ".\manifold.py", line 54, in <module>
#     x_manifold = mdl.inverse(z_manifold)
# AttributeError: 'SimpleRealNVP' object has no attribute 'inverse'

What is the appropriate way to include extended methods/properties into the SavedModel?

(I've also tried specifying custom_objects in load_model() but this didn't work.)