Problem with enumeration for batch data in Deep markov models

Thanks,
I’ve played with the batch dimension, as it made me go further with the debugging. Here I try the batch_dim=-1 and get the error - ValueError: at site "emit_x_0", invalid log_prob shape.
However, all hmm examples use the dim=-2, I guess because they don’t want to be constrained by 1D output dimension. Accordingly, I also tried a dim=-2 version with no success.

Bottom line - using dim=-1 and removing the obs_plate allowed me to move further in debugging to the svi.step (with the aforementioned format_shapes), but failed on the same “double-batch” shape assertion.

I will provide a minimal runnable example today or tomorrow