Hi @BigDiaos, I think the issue here is that Bernoulli and Categorical accept parameters of different shape:
-
Bernoulliaccepts one parameter per instance (per “coin” flipped), whereas -
Categoricalaccepts D-many parameters per instance of choice among D categories.
I think if you want to update the tutorial code, you’ll need to add an extra shape (num_categories,) to the right of your loc parameters, requiring changing shapes of the encoder and decoder networks.