leras: add ability to save sub layers in a dict

2020-07-01 22:17:35 +04:00
parent e0a1d52d78
commit e6aa996814
1 changed files with 6 additions and 2 deletions
@@ -18,6 +18,10 @@ class ModelBase(nn.Saveable):
        if isinstance (layer, list):
            for i,sublayer in enumerate(layer):
                self._build_sub(sublayer, f"{name}_{i}")
+        elif isinstance (layer, dict):
+            for subname in layer.keys():
+                sublayer = layer[subname]
+                self._build_sub(sublayer, f"{name}_{subname}")
        elif isinstance (layer, nn.LayerBase) or \
                isinstance (layer, ModelBase):