feat: add mean and logvar info to variational bottleneck

Author: flavioschneider

audio_diffusion_pytorch/modules.py

@@ -1211,7 +1211,7 @@ def forward(
         logvar = torch.clamp(logvar, -30.0, 20.0)
         out = gaussian_sample(mean, logvar)
         loss = kl_loss(mean, logvar) * self.loss_weight
-        return (out, dict(loss=loss)) if with_info else out
+        return (out, dict(loss=loss, mean=mean, logvar=logvar)) if with_info else out
 
 
 class AutoEncoder1d(nn.Module):

setup.py

@@ -3,7 +3,7 @@
 setup(
     name="audio-diffusion-pytorch",
     packages=find_packages(exclude=[]),
-    version="0.0.64",
+    version="0.0.65",
     license="MIT",
     description="Audio Diffusion - PyTorch",
     long_description_content_type="text/markdown",