如何在张量流概率分布中保存和加载训练变量？

Question

我想建立一个流模型，可以保存经过训练的模型。

flow = tfd.TransformedDistribution(
    ...
    distribution=tfd.Normal(loc=0.0, scale=1.0),
    bijector=my_bijectors
)

for e in epochs:
    ...
    with tf.GradientTape() as tape:
         log_prob_loss = loss() # some loss func here
    grads = tape.gradient(log_prob_loss, flow.trainable_variables)
    optimizer.apply_gradients(zip(grads, flow.trainable_variables))
    ...

如何在此处保存和加载flow？

我当时正在考虑使用泡菜保存trainable_variables，但是我不知道如何将已加载的经过训练的trainable_variables应用于原始的。有办法更新吗？