x-transformers 2.3.9__tar.gz → 2.3.10__tar.gz

{x_transformers-2.3.9 → x_transformers-2.3.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.3.9
+Version: 2.3.10
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers

{x_transformers-2.3.9 → x_transformers-2.3.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-transformers"
-version = "2.3.9"
+version = "2.3.10"
 description = "X-Transformers"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{x_transformers-2.3.9 → x_transformers-2.3.10}/x_transformers/continuous.py

@@ -42,10 +42,9 @@ def masked_mean(t, mask):
 # probabilistic loss fn
 
 class GaussianNLL(Module):
-    def forward(self, pred, target):
+    def forward(self, pred, target, eps = 1e-5):
         mean, var = pred
-        dist = Normal(mean, var)
-        return -dist.log_prob(target)
+        return F.gaussian_nll_loss(mean, target, var)
 
 # main classes
 
@@ -271,7 +270,9 @@ class ContinuousAutoregressiveWrapper(Module):
 
             if self.probabilistic:
                 mean, var = last_output
-                last_output = torch.normal(mean, var * temperature)
+                stddev = var.clamp(min = 1e-5).sqrt()
+
+                last_output = torch.normal(mean, stddev * temperature)
 
             out = cat((out, last_output), dim = -2)