diff --git a/src/easyvae/autoencoder.py b/src/easyvae/autoencoder.py
index 1dc6c24..1880781 100644
--- a/src/easyvae/autoencoder.py
+++ b/src/easyvae/autoencoder.py
@@ -232,7 +232,7 @@ class VariationalAutoencoder(AAutoencoder):
         code = self.encoder.forward(v)
         sample = self.sampler.forward(code)
         out = self.decoder.forward(sample)
-        return out, code
+        return out, sample
 
     def encode(self, v: np.ndarray) -> np.ndarray:
         return self.sampler.forward(