Fix missing config param in attention layers

taylorhansen · Sep 24, 2023 · a411263 · a411263
1 parent c7cf3c7
commit a411263
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/src/py/models/utils/attention.py b/src/py/models/utils/attention.py
@@ -89,6 +89,7 @@ def __init__(
         super().__init__(**kwargs)
         self.num_heads = num_heads
         self.depth = depth
+        self.use_bias = use_bias
 
         dim = num_heads * depth
         self.query = tf.keras.layers.Dense(
@@ -234,7 +235,11 @@ def get_config(self):
         return (
             super().get_config()
             | self.get_initializer_config()
-            | {"num_heads": self.num_heads, "depth": self.depth}
+            | {
+                "num_heads": self.num_heads,
+                "depth": self.depth,
+                "use_bias": self.use_bias,
+            }
         )
 
 
@@ -331,6 +336,7 @@ def get_config(self):
                 "depth": self.mha.depth,
                 "rff": self.rff,
                 "use_layer_norm": self.use_layer_norm,
+                "use_bias": self.mha.use_bias,
             }
         )
 
@@ -409,6 +415,7 @@ def get_config(self):
                 "depth": self.mab.mha.depth,
                 "rff": self.mab.rff,
                 "use_layer_norm": self.mab.use_layer_norm,
+                "use_bias": self.mab.mha.use_bias,
             }
         )
 
@@ -514,6 +521,7 @@ def get_config(self):
                 "rff": self.mab.rff,
                 "rff_s": self.rff_s,
                 "use_layer_norm": self.mab.use_layer_norm,
+                "use_bias": self.mab.mha.use_bias,
                 "seed_initializer": tf.keras.initializers.serialize(
                     self.seed_initializer
                 ),