Many updates, cool stuff on the way.

1 files changed, 7 insertions, 2 deletions

diff --git a/src/text_recognizer/networks/util.py b/src/text_recognizer/networks/util.py
index 711a952..131a6b4 100644
--- a/src/text_recognizer/networks/util.py
+++ b/src/text_recognizer/networks/util.py
@@ -65,13 +65,18 @@ def configure_backbone(backbone: str, backbone_args: Dict) -> Type[nn.Module]:
         network_args = state_dict["network_args"]
         weights = state_dict["model_state"]
 
+        freeze = False
+        if "freeze" in backbone_args and backbone_args["freeze"] is True:
+            backbone_args.pop("freeze")
+            freeze = True
+        network_args = backbone_args
+
         # Initializes the network with trained weights.
         backbone = backbone_(**network_args)
         backbone.load_state_dict(weights)
-        if "freeze" in backbone_args and backbone_args["freeze"] is True:
+        if freeze:
             for params in backbone.parameters():
                 params.requires_grad = False
-
     else:
         backbone_ = getattr(network_module, backbone)
         backbone = backbone_(**backbone_args)