fix a bug thanks to @almsley

lucidrains · Feb 15, 2021 · 4d3da00 · 4d3da00
1 parent 3531538
commit 4d3da00
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/big_sleep/big_sleep.py b/big_sleep/big_sleep.py
@@ -208,7 +208,8 @@ def forward(self, text_embed, return_loss = True):
             skews = torch.mean(torch.pow(zscores, 3.0))
             kurtoses = torch.mean(torch.pow(zscores, 4.0)) - 3.0
 
-        lat_loss = lat_loss + torch.abs(kurtoses) / num_latents + torch.abs(skews) / num_latents
+            lat_loss = lat_loss + torch.abs(kurtoses) / num_latents + torch.abs(skews) / num_latents
+
         cls_loss = ((50 * torch.topk(soft_one_hot_classes, largest = False, dim = 1, k = 999)[0]) ** 2).mean()
 
         sim_loss = -self.loss_coef * torch.cosine_similarity(text_embed, image_embed, dim = -1).mean()

diff --git a/big_sleep/version.py b/big_sleep/version.py
@@ -1 +1 @@
-__version__ = '0.5.0'
+__version__ = '0.5.1'