dist.Bernoulli recognize the right batch_size, but dist.Categorical won't

Hi @BigDiaos, I think the issue here is that Bernoulli and Categorical accept parameters of different shape:

  • Bernoulli accepts one parameter per instance (per “coin” flipped), whereas
  • Categorical accepts 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.