From e6c1229a7e422acf760a374f854d1a56899c6669 Mon Sep 17 00:00:00 2001
From: Lenoctambule <106790775+lenoctambule@users.noreply.github.com>
Date: Mon, 13 Apr 2026 23:07:11 +0200
Subject: [PATCH] fix: return encoding after sample layer

---
 src/easyvae/autoencoder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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(