titans-pytorch 0.0.8__tar.gz → 0.0.9__tar.gz

{titans_pytorch-0.0.8 → titans_pytorch-0.0.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.8
+Version: 0.0.9
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -77,6 +77,18 @@ retrieved = mem(seq)
 assert seq.shape == retrieved.shape
 ```
 
+## Experiments
+
+```bash
+$ pip install .[examples]
+```
+
+Then
+
+```bash
+$ python train.py
+```
+
 ## Citations
 
 ```bibtex

{titans_pytorch-0.0.8 → titans_pytorch-0.0.9}/README.md

@@ -30,6 +30,18 @@ retrieved = mem(seq)
 assert seq.shape == retrieved.shape
 ```
 
+## Experiments
+
+```bash
+$ pip install .[examples]
+```
+
+Then
+
+```bash
+$ python train.py
+```
+
 ## Citations
 
 ```bibtex

{titans_pytorch-0.0.8 → titans_pytorch-0.0.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.0.8"
+version = "0.0.9"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{titans_pytorch-0.0.8 → titans_pytorch-0.0.9}/titans_pytorch/titans.py

@@ -311,7 +311,10 @@ class NeuralMemory(Module):
         past_state: tuple[dict[str, Tensor], dict[str, Tensor]] | None = None,
         return_next_memories = False
     ):
-        batch = seq.shape[0]
+        batch, seq_len = seq.shape[:2]
+
+        if seq_len < self.chunk_size:
+            return torch.zeros_like(seq)
 
         if exists(past_state):
             past_state = tuple(TensorDict(d) for d in past_state)