Repository: pyro-ppl/pyro

Issue:
vae comparison example giving different results for Pyro and PyTorch implementations
Hi,

Pyro version: 3.3
PyTorch: 1.1

Running `examples/vae/vae_comparison.py`

With PyTorch:

```
downloading data
download complete.
downloading data
download complete.
Running PyTorch VAE implementation
====> Epoch: 0 
Training loss: 0.0022
Test set loss: 0.0027
====> Epoch: 1 
Training loss: 0.0021
Test set loss: 0.0024
====> Epoch: 2 
Training loss: 0.0021
Test set loss: 0.0022
====> Epoch: 3 
Training loss: 0.0021
Test set loss: 0.0021
====> Epoch: 4 
Training loss: 0.0021
Test set loss: 0.0021
====> Epoch: 5 
Training loss: 0.0021
Test set loss: 0.0021
====> Epoch: 6 
Training loss: 0.0021
Test set loss: 0.0021
====> Epoch: 7 
Training loss: 0.0021
Test set loss: 0.0021
====> Epoch: 8 
Training loss: 0.0021
Test set loss: 0.0021
====> Epoch: 9 
Training loss: 0.0021
Test set loss: 0.0021
```

With Pyro:

```
downloading data
download complete.
downloading data
download complete.
Running Pyro VAE implementation
====> Epoch: 0 
Training loss: 0.0017
Test set loss: 0.0013
====> Epoch: 1 
Training loss: 0.0012
Test set loss: 0.0012
====> Epoch: 2 
Training loss: 0.0011
Test set loss: 0.0011
====> Epoch: 3 
Training loss: 0.0011
Test set loss: 0.0011
====> Epoch: 4 
Training loss: 0.0011
Test set loss: 0.0011
====> Epoch: 5 
Training loss: 0.0011
Test set loss: 0.0011
====> Epoch: 6 
Training loss: 0.0011
Test set loss: 0.0011
====> Epoch: 7 
Training loss: 0.0011
Test set loss: 0.0011
====> Epoch: 8 
Training loss: 0.0011
Test set loss: 0.0011
====> Epoch: 9 
Training loss: 0.0011
Test set loss: 0.0011
```

Is this difference expected?


Discussion:
@neerajprad shouldn't it be `.rsample()` here: https://github.com/pyro-ppl/pyro/blob/dev/examples/vae/vae_comparison.py#L106

Changing that line gives the following:

```
downloading data
download complete.
downloading data
download complete.
Running PyTorch VAE implementation
====> Epoch: 0 
Training loss: 0.0016
Test set loss: 0.0012
====> Epoch: 1 
Training loss: 0.0012
Test set loss: 0.0011
====> Epoch: 2 
Training loss: 0.0011
Test set loss: 0.0010
====> Epoch: 3 
Training loss: 0.0011
Test set loss: 0.0010
====> Epoch: 4 
Training loss: 0.0011
Test set loss: 0.0010
====> Epoch: 5 
Training loss: 0.0011
Test set loss: 0.0010
====> Epoch: 6 
Training loss: 0.0011
Test set loss: 0.0010
```
Good catch! That is being called from `compute_loss_and_gradient` in the PyTorch implementation. I was under the impression that its only being used during evaluation. Do you mind submitting your fix in a PR? 🙂

Resolving patch:
diff --git a/examples/vae/vae_comparison.py b/examples/vae/vae_comparison.py
--- a/examples/vae/vae_comparison.py
+++ b/examples/vae/vae_comparison.py
@@ -103,7 +103,7 @@ def model_eval(self, x):
         """
         z_mean, z_var = self.vae_encoder(x)
         if self.mode == TRAIN:
-            z = Normal(z_mean, z_var.sqrt()).sample()
+            z = Normal(z_mean, z_var.sqrt()).rsample()
         else:
             z = z_mean
         return self.vae_decoder(z), z_mean, z_var


Test patch: