Repository: pyro-ppl/pyro

Issue:
another error in ss_vae_M2.py
As running ss_vae_M2.py in the latest pyro and latest pytorch, I have another error if `--aux-loss` is on:
```
Traceback (most recent call last):
  File "/home/jbaik/.pyenv/versions/3.6.4/lib/python3.6/site-packages/pyro_ppl-0.1.2-py3.6.egg/pyro/poutine/trace.py", line 81, in log_pdf
    site_log_p = site["log_pdf"]
KeyError: 'log_pdf'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "ss_vae_M2.py", line 453, in <module>
    run_inference_ss_vae(args)
  File "ss_vae_M2.py", line 360, in run_inference_ss_vae
    run_inference_for_epoch(data_loaders, losses, periodic_interval_batches)
  File "ss_vae_M2.py", line 254, in run_inference_for_epoch
    new_loss = losses[loss_id].step(xs, ys)
  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 95, in _get_traces
    log_r = model_trace.log_pdf() - guide_trace.log_pdf()
  File "/home/jbaik/.pyenv/versions/3.6.4/lib/python3.6/site-packages/pyro_ppl-0.1.2-py3.6.egg/pyro/poutine/trace.py", line 85, in log_pdf
    site["value"], *args, **kwargs).sum() * site["scale"]
  File "/home/jbaik/.pyenv/versions/3.6.4/lib/python3.6/site-packages/pyro_ppl-0.1.2-py3.6.egg/pyro/distributions/random_primitive.py", line 52, in log_prob
    return self.dist_class(*args, **kwargs).log_prob(x)
  File "/home/jbaik/.pyenv/versions/3.6.4/lib/python3.6/site-packages/pyro_ppl-0.1.2-py3.6.egg/pyro/distributions/torch_wrapper.py", line 48, in log_prob
    assert len(self.log_pdf_mask.size()) <= len(log_prob.size())
AttributeError: 'int' object has no attribute 'size'
```

Could you take a look for this too? I couldn't reopen the closed issue #725 so have opened a new issue here.


Discussion:
Thanks for the bug report, I'll take a look later today.

Resolving patch:
diff --git a/examples/ss_vae_M2.py b/examples/ss_vae_M2.py
--- a/examples/ss_vae_M2.py
+++ b/examples/ss_vae_M2.py
@@ -188,11 +188,8 @@ def model_classify(self, xs, ys=None):
             # similar to the NIPS 14 paper (Kingma et al).
             if ys is not None:
                 alpha = self.encoder_y.forward(xs)
-                pyro.sample("y_aux",
-                            dist.one_hot_categorical,
-                            alpha,
-                            log_pdf_mask=self.aux_loss_multiplier,
-                            obs=ys)
+                with pyro.poutine.scale(None, self.aux_loss_multiplier):
+                    pyro.sample("y_aux", dist.one_hot_categorical, alpha, obs=ys)
 
     def guide_classify(self, xs, ys=None):
         """


Test patch:
diff --git a/tests/test_examples.py b/tests/test_examples.py
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -31,8 +31,9 @@ def discover_examples():
                 continue
             example = os.path.relpath(path, EXAMPLES_DIR)
             CPU_EXAMPLES.append((example, args))
-            if '--enum-discrete' in text:
-                CPU_EXAMPLES.append((example, args + ['--enum-discrete']))
+            for flag in ['--enum-discrete', '--aux-loss']:
+                if flag in text:
+                    CPU_EXAMPLES.append((example, args + [flag]))
             if '--cuda' in text:
                 CUDA_EXAMPLES.append((example, args + ['--cuda']))
     CPU_EXAMPLES.sort()