From d460e088760af0352f979daa29fd0e7b96708919 Mon Sep 17 00:00:00 2001
From: Radev <radevs@win.rpi.edu>
Date: Tue, 14 May 2024 18:49:15 -0400
Subject: [PATCH] Add config

---
 bayesflow/experimental/networks/resnet/hidden_block.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/bayesflow/experimental/networks/resnet/hidden_block.py b/bayesflow/experimental/networks/resnet/hidden_block.py
index c33bc94a..42305209 100644
--- 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