Add config

bayesflow-org · May 14, 2024 · d460e08 · d460e08
1 parent faae8df
commit d460e08
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/bayesflow/experimental/networks/resnet/hidden_block.py b/bayesflow/experimental/networks/resnet/hidden_block.py
@@ -34,8 +34,9 @@ def call(self, inputs, training=False):
     def get_config(self):
         config = super().get_config()
         config.update({
-            "units": self.activation_fn,
             "residual": self.residual,
-            "spectral_norm": self.spectral_norm
+            "spectral_norm": self.spectral_norm,
+            "activation_fn": keras.saving.serialize_keras_object(self.activation_fn),
+            "dense_with_dropout": keras.saving.serialize_keras_object(self.dense_with_dropout)
         })
         return config