x-transformers 2.1.27__tar.gz → 2.1.28__tar.gz

{x_transformers-2.1.27 → x_transformers-2.1.28}/PKG-INFO

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

{x_transformers-2.1.27 → x_transformers-2.1.28}/pyproject.toml

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

{x_transformers-2.1.27 → x_transformers-2.1.28}/x_transformers/belief_state_wrapper.py

@@ -260,6 +260,7 @@ class BeliefStateWrapper(Module):
 
         if exists(lens):
             mask = einx.less('j, i -> i j', arange(seq_len, device = device), lens)
+            seq_for_labels = torch.where(mask, seq, -1)
 
         # forward autoregressive
 
@@ -319,7 +320,7 @@ class BeliefStateWrapper(Module):
 
         labels_fi, labels_bi = (fi + 1), (bi - 1)
 
-        forward_labels, backward_labels = seq[:, labels_fi], seq[:, labels_bi]
+        forward_labels, backward_labels = seq_for_labels[:, labels_fi], seq_for_labels[:, labels_bi]
 
         labels = cat((forward_labels, backward_labels), dim = -1)
 
@@ -337,7 +338,8 @@ class BeliefStateWrapper(Module):
         loss = F.cross_entropy(
             rearrange(logits, 'b n (fb l) -> b l (fb n)', fb = 2),
             labels,
-            reduction = 'none' if self.needs_loss_weight else 'mean'
+            reduction = 'none' if self.needs_loss_weight else 'mean',
+            ignore_index = -1
         )
 
         # maybe predict terminal