Saving transform_distribution with tf.train.checkpoint encounter 'Unable to save the object ListWrapper' error

ValueError: Unable to save the object ListWrapper([16, 16, 4]) (a list wrapper constructed to track trackable TensorFlow objects). The wrapped list was modified outside the wrapper (its final value was [16, 16, 4], its value when a checkpoint dependency was added was None), which breaks restoration on object creation.

If you don't need this list checkpointed, wrap it in a tf.contrib.checkpoint.NoDependency object; it will be automatically un-wrapped and subsequently ignored.

Everything went fine until I define a bijector myself then add it to the chain. The bijector was like that:

class Bij(tfb.Bijector):
    def __init___(...):
        super().__init__(..., forward_min_event_ndims=3)
        self.bi1 = Bijector1()
        self.bi2 = Bijector2()

    def _forward(self, x):
        x = self.bi1.forward(x)
        # some ops on x here
        return self.bi2.forward(x)
        # I hope the way it forwards is right, I can't just chain them up cuz some ops on x between bi1 and bi2. If there's better way to achieve this, pls let me know!

    ...
    
    def _forward_log_det_jacobian(self, x):
        # I use bi1.forward_log_det_jacobian(x) to get
        # jaconbian then summed it with bi2.forward_log_det_jacobian(x) with parameter
        # event_ndims = forward_min_event_ndims in both forward_log_det_jacobian().

The training goes fine but when I was trying to save it at the end of the epoch, it throws this error.

Start by examining the custom Bij class and the tf.train.checkpoint save operation described in the report. Reproduce the failure with the transform_distribution chain, then determine which tracked list is modified and define done as a checkpoint that saves successfully and restores the custom bijector.