Merge pull request #24 from Pasewark/token_inference_fix

Small change so token embeddings aren't looked up for past tokens during inference

Author: lucidrains

native_sparse_attention_pytorch/transformer.py

@@ -273,7 +273,10 @@ def forward(
 
         # token embedding
 
-        tokens = self.token_emb(ids)
+        if is_inferencing:
+            tokens = self.token_emb(ids[:, -1:])
+        else:
+            tokens = self.token_emb(ids)
 
         # prepare maybe flex attention masks
 
@@ -298,9 +301,6 @@ def forward(
 
         next_cache = []
 
-        if is_inferencing:
-            tokens = tokens[:, -1:]
-
         # layers
 
         for attn, ff in self.layers: