(Reproducible example) Proper shape for Predictive with parallel=True

Anyone have thoughts on this? I’m starting to think a bug is occurring with the dimensions of weights. @eb8680_2 previously wrote:

e.g. Trace_ELBO(num_particles=100, vectorize_particles=True)` ; the ELBO implementation will automatically create an extra batch dimension in your model and guide and average over this dimension when computing the ELBO estimate.

However, here weight has been given two additional dimensions, not one. Same for bias.

I am still trying to track down the causes for this but I seem to be stuck.