gate the value heads in taylor linear attention

lucidrains · Jan 15, 2024 · 6c70a6e · 6c70a6e
1 parent 298a0db
commit 6c70a6e
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/equiformer_pytorch/equiformer_pytorch.py b/equiformer_pytorch/equiformer_pytorch.py
@@ -588,7 +588,7 @@ def __init__(
 
         if self.has_linear_attn:
             degree_zero_dim = fiber[0]
-            self.linear_attn = TaylorSeriesLinearAttn(degree_zero_dim, dim_head = linear_attn_dim_head, heads = num_linear_attn_heads, combine_heads = False)
+            self.linear_attn = TaylorSeriesLinearAttn(degree_zero_dim, dim_head = linear_attn_dim_head, heads = num_linear_attn_heads, combine_heads = False, gate_value_heads = True)
             hidden_fiber = tuple_set_at_index(hidden_fiber, 0, hidden_fiber[0] + linear_attn_dim_head * num_linear_attn_heads)
 
         # gating heads across all degree outputs

diff --git a/equiformer_pytorch/version.py b/equiformer_pytorch/version.py
@@ -1 +1 @@
-__version__ = '0.5.0'
+__version__ = '0.5.1'
diff --git a/setup.py b/setup.py
@@ -26,7 +26,7 @@
     'einops>=0.6',
     'filelock',
     'opt-einsum',
-    'taylor-series-linear-attention>=0.0.11',
+    'taylor-series-linear-attention>=0.1.4',
     'torch>=1.6',
   ],
   setup_requires=[