Getting an error in svi step due to to a multiclass distribution in sample using pyro and pytorch

Question

I\'m working on a causal variational autoencoder which works with class segmentation masks, class labels and causality(0 or 1) as the inputs.

I\'m getting an error when w