Fix loading serialized transformers

Author: stefanradev93

bayesflow/networks/transformers/fusion_transformer.py

@@ -4,6 +4,7 @@
 
 from bayesflow.types import Tensor
 from bayesflow.utils import check_lengths_same
+from bayesflow.utils.decorators import sanitize_input_shape
 
 from ..summary_network import SummaryNetwork
 
@@ -151,3 +152,8 @@ def call(self, input_sequence: Tensor, training: bool = False, **kwargs) -> Tens
         summary = self.attention_blocks[-1](keras.ops.expand_dims(template, axis=1), rep, training=training, **kwargs)
         summary = self.output_projector(keras.ops.squeeze(summary, axis=1))
         return summary
+
+    @sanitize_input_shape
+    def build(self, input_shape):
+        super().build(input_shape)
+        self.call(keras.ops.zeros(input_shape))

bayesflow/networks/transformers/set_transformer.py

@@ -3,6 +3,7 @@
 
 from bayesflow.types import Tensor
 from bayesflow.utils import check_lengths_same
+from bayesflow.utils.decorators import sanitize_input_shape
 
 from ..summary_network import SummaryNetwork
 
@@ -150,3 +151,8 @@ def call(self, input_set: Tensor, training: bool = False, **kwargs) -> Tensor:
         summary = self.pooling_by_attention(summary, training=training, **kwargs)
         summary = self.output_projector(summary)
         return summary
+
+    @sanitize_input_shape
+    def build(self, input_shape):
+        super().build(input_shape)
+        self.call(keras.ops.zeros(input_shape))

bayesflow/networks/transformers/time_series_transformer.py

@@ -3,6 +3,7 @@
 
 from bayesflow.types import Tensor
 from bayesflow.utils import check_lengths_same
+from bayesflow.utils.decorators import sanitize_input_shape
 
 from ..embeddings import Time2Vec, RecurrentEmbedding
 from ..summary_network import SummaryNetwork
@@ -147,3 +148,8 @@ def call(self, input_sequence: Tensor, training: bool = False, **kwargs) -> Tens
         summary = self.pooling(inp)
         summary = self.output_projector(summary)
         return summary
+
+    @sanitize_input_shape
+    def build(self, input_shape):
+        super().build(input_shape)
+        self.call(keras.ops.zeros(input_shape))