feat: add mean and logvar info to variational bottleneck

archinetai · Oct 15, 2022 · 58db626 · 58db626
1 parent 7fabeb3
commit 58db626
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/audio_diffusion_pytorch/modules.py b/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):

diff --git a/setup.py b/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",