last tweak

lucidrains · Jan 20, 2025 · ec05100 · ec05100
1 parent 28eddd6
commit ec05100
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.1.8"
+version = "0.1.9"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

diff --git a/titans_pytorch/titans.py b/titans_pytorch/titans.py
@@ -391,6 +391,8 @@ def forward_and_loss(params, inputs, loss_weights, target):
 
         # whether to use averaging of chunks, or attention pooling
 
+        assert not (attn_pool_chunks and chunk_size == 1), '`attn_pool_chunks` cannot be set to True if `chunk_size` is set to 1'
+
         if not attn_pool_chunks:
             chunk_reduce_module = Reduce('b (n c) ... -> b n ...', 'mean', c = chunk_size)
         else: