dot_prod needs to unpack arguments from tuple

qarac/models/layers/GlobalAttentionPoolingHead.py

@@ -11,7 +11,8 @@ import tensorflow
 
 
 @tensorflow.function
-def dot_prod(x,y):
+def dot_prod(vectors):
+    (x,y) = vectors
     return tensorflow.tensordot(x,y,axes=1)
     
 
@@ -81,6 +82,6 @@ class GlobalAttentionPoolingHead(keras.layers.Layer):
         lp = tensorflow.linalg.l2_normalize(tensorflow.vectorized_map(self.project_local,
                                                                       X),
                                             axis=2)
-        attention = tensorflow.vectorized_map(dot_prod,[lp,gp])
+        attention = tensorflow.vectorized_map(dot_prod,(lp,gp))
         return tensorflow.reduce_sum(attention *X,
                                      axis=1)