Update SelfAttention_Family.py

Author: wuhaixu2016

layers/SelfAttention_Family.py

@@ -143,7 +143,7 @@ def _update_context(self, context_in, V, scores, index, L_Q, attn_mask):
         else:
             return (context_in, None)
 
-    def forward(self, queries, keys, values, attn_mask):
+    def forward(self, queries, keys, values, attn_mask, tau=None, delta=None):
         B, L_Q, H, D = queries.shape
         _, L_K, _, _ = keys.shape
 
@@ -203,7 +203,9 @@ def forward(self, queries, keys, values, attn_mask, tau=None, delta=None):
             queries,
             keys,
             values,
-            attn_mask
+            attn_mask,
+            tau=tau,
+            delta=delta
         )
         out = out.view(B, L, -1)