ready for use in another project

Author: lucidrains

agent_attention_pytorch/agent_transformer.py

@@ -65,7 +65,7 @@ def __init__(
             nn.Sigmoid()
         ) if gate else None
 
-        self.qa_talking_heads = nn.Conv2d(heads, heads, 1, bias = False) if talking_heads else nn.Identity()
+        self.talking_heads = nn.Conv2d(heads, heads, 1, bias = False) if talking_heads else nn.Identity()
         self.ak_talking_heads = nn.Conv2d(heads, heads, 1, bias = False) if talking_heads else nn.Identity()
 
         self.qa_dropout = nn.Dropout(dropout)
@@ -146,6 +146,7 @@ def __init__(
         heads = 8,
         dim_head = 64,
         ff_mult = 4,
+        final_norm = True,
         **attn_kwargs: dict
     ):
         super().__init__()
@@ -167,6 +168,8 @@ def __init__(
                 FeedForward(dim = dim, mult = ff_mult)
             ]))
 
+        self.final_norm = RMSNorm(dim) if final_norm else None
+
     def forward(
         self,
         x,
@@ -192,6 +195,10 @@ def forward(
 
             a, x = unpack(x, ps, 'b * d')
 
+        if exists(self.final_norm):
+            x = self.final_norm(x)
+            a = self.final_norm(a)
+
         if not return_agent_tokens:
             return x
 

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'agent-attention-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.1.5',
+  version = '0.1.6',
   license='MIT',
   description = 'Agent Attention - Pytorch',
   author = 'Phil Wang',