Unhashable Type TypeError: Tensors are unhashable. Instead, use tensor.ref() as the key

Question

I am trying to implement a custom variational autoencoder. Following is the code to reproduce.

epsilon_std = 1.0

vx = tf.keras.layers.Input(batch_shape=(None, ma