Bring back MultiHeadAttention after keras bugfix

Author: stefanradev93

bayesflow/networks/transformers/mab.py

@@ -4,7 +4,6 @@
 
 from bayesflow.types import Tensor
 from bayesflow.networks import MLP
-from .mha import MultiHeadAttention
 
 
 @serializable(package="bayesflow.networks")
@@ -40,8 +39,8 @@ def __init__(
         super().__init__(**kwargs)
 
         self.input_projector = layers.Dense(embed_dim)
-        self.attention = MultiHeadAttention(
-            embed_dim=embed_dim,
+        self.attention = layers.MultiHeadAttention(
+            key_dim=embed_dim,
             num_heads=num_heads,
             dropout=dropout,
             use_bias=use_bias,