# Buggy code
def model(self, xs, ys=None):
    """
    The model corresponds to the following generative process:
    p(z) = normal(0,I)              # handwriting style (latent)
    p(y|x) = categorical(I/10.)     # which digit (semi-supervised)
    p(x|y,z) = bernoulli(mu(y,z))   # an image
    mu is given by a neural network  `decoder`

    :param xs: a batch of scaled vectors of pixels from an image
    :param ys: (optional) a batch of the class labels i.e.
               the digit corresponding to the image(s)
    :return: None
    """
    # register this pytorch module and all of its sub-modules with pyro
    pyro.module("ss_vae", self)

    batch_size = xs.size(0)
    with pyro.iarange("independent"):
        # sample the handwriting style from the constant prior distribution
        prior_mu = Variable(torch.zeros([batch_size, self.z_dim]))
        prior_sigma = Variable(torch.ones([batch_size, self.z_dim]))
        zs = pyro.sample("z", dist.normal, prior_mu, prior_sigma)

        # if the label y (which digit to write) is supervised, sample from the
        # constant prior, otherwise, observe the value (i.e. score it against the constant prior)
        alpha_prior = Variable(
            torch.ones([batch_size, self.output_size]) / (1.0 * self.output_size)
        )
        if ys is None:
            ys = pyro.sample("y", dist.one_hot_categorical, alpha_prior)
        else:
            pyro.sample("y", dist.one_hot_categorical, alpha_prior, obs=ys)

        # finally, score the image (x) using the handwriting style (z) and
        # the class label y (which digit to write) against the
        # parametrized distribution p(x|y,z) = bernoulli(decoder(y,z))
        # where `decoder` is a neural network
        mu = self.decoder.forward([zs, ys])
        pyro.sample("x", dist.bernoulli, mu, obs=xs)


# Traceback (NotImplementedError)
$ python ss_vae_M2.py -enum
Traceback (most recent call last):
File "ss_vae_M2.py", line 464, in <module>
run_inference_ss_vae(args)
File "ss_vae_M2.py", line 371, in run_inference_ss_vae
run_inference_for_epoch(data_loaders, losses, periodic_interval_batches)
File "ss_vae_M2.py", line 268, in run_inference_for_epoch
new_loss = losses[loss_id].step(xs)
File "/home/jbaik/.pyenv/versions/3.6.4/lib/python3.6/site-packages/pyro_ppl-0.1.2-py3.6.egg/pyro/infer/svi.py", line 98, in step
loss = self.loss_and_grads(self.model, self.guide, *args, **kwargs)
File "/home/jbaik/.pyenv/versions/3.6.4/lib/python3.6/site-packages/pyro_ppl-0.1.2-py3.6.egg/pyro/infer/elbo.py", line 65, in loss_and_grads
return self.which_elbo.loss_and_grads(model, guide, *args, **kwargs)
File "/home/jbaik/.pyenv/versions/3.6.4/lib/python3.6/site-packages/pyro_ppl-0.1.2-py3.6.egg/pyro/infer/trace_elbo.py", line 146, in loss_and_grads
for weight, model_trace, guide_trace, log_r in self._get_traces(model, guide, *args, **kwargs):
File "/home/jbaik/.pyenv/versions/3.6.4/lib/python3.6/site-packages/pyro_ppl-0.1.2-py3.6.egg/pyro/infer/trace_elbo.py", line 79, in _get_traces
check_enum_discrete_can_run(model_trace, guide_trace)
File "/home/jbaik/.pyenv/versions/3.6.4/lib/python3.6/site-packages/pyro_ppl-0.1.2-py3.6.egg/pyro/infer/trace_elbo.py", line 45, in check_enum_discrete_can_run
for shape, (source, name) in sorted(shapes.items())])))
NotImplementedError: enum_discrete does not support mixture of batched and un-batched variables. Try rewriting your model to avoid batching or running with enum_discrete=False. Found the following variables of different batch shapes:
guide y: shape = (200,)
guide z: shape = (200, 50)
model x: shape = (200, 784)

# Fixed code
def model(self, xs, ys=None):
    """
    The model corresponds to the following generative process:
    p(z) = normal(0,I)              # handwriting style (latent)
    p(y|x) = categorical(I/10.)     # which digit (semi-supervised)
    p(x|y,z) = bernoulli(mu(y,z))   # an image
    mu is given by a neural network  `decoder`

    :param xs: a batch of scaled vectors of pixels from an image
    :param ys: (optional) a batch of the class labels i.e.
               the digit corresponding to the image(s)
    :return: None
    """
    # register this pytorch module and all of its sub-modules with pyro
    pyro.module("ss_vae", self)

    batch_size = xs.size(0)
    with pyro.iarange("independent"):
        # sample the handwriting style from the constant prior distribution
        prior_mu = Variable(torch.zeros([batch_size, self.z_dim]))
        prior_sigma = Variable(torch.ones([batch_size, self.z_dim]))
        zs = pyro.sample("z", dist.normal, prior_mu, prior_sigma, extra_event_dims=1)

        # if the label y (which digit to write) is supervised, sample from the
        # constant prior, otherwise, observe the value (i.e. score it against the constant prior)
        alpha_prior = Variable(
            torch.ones([batch_size, self.output_size]) / (1.0 * self.output_size)
        )
        if ys is None:
            ys = pyro.sample("y", dist.one_hot_categorical, alpha_prior)
        else:
            pyro.sample("y", dist.one_hot_categorical, alpha_prior, obs=ys)

        # finally, score the image (x) using the handwriting style (z) and
        # the class label y (which digit to write) against the
        # parametrized distribution p(x|y,z) = bernoulli(decoder(y,z))
        # where `decoder` is a neural network
        mu = self.decoder.forward([zs, ys])
        pyro.sample("x", dist.bernoulli, mu, extra_event_dims=1, obs=xs)