add compute_output_shape to mlp

bayesflow-org · Jul 11, 2024 · 4b629e4 · 4b629e4
1 parent d8428c8
commit 4b629e4
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/bayesflow/networks/mlp/mlp.py b/bayesflow/networks/mlp/mlp.py
@@ -81,3 +81,9 @@ def call(self, x: Tensor, **kwargs) -> Tensor:
         for layer in self.res_blocks:
             x = layer(x, training=kwargs.get("training", False))
         return x
+
+    def compute_output_shape(self, input_shape):
+        for layer in self.res_blocks:
+            input_shape = layer.compute_output_shape(input_shape)
+
+        return input_shape