leras: add ability to save sub layers in a dict

This commit is contained in:
Colombo 2020-07-01 22:17:35 +04:00
parent e0a1d52d78
commit e6aa996814

View file

@ -18,6 +18,10 @@ class ModelBase(nn.Saveable):
if isinstance (layer, list): if isinstance (layer, list):
for i,sublayer in enumerate(layer): for i,sublayer in enumerate(layer):
self._build_sub(sublayer, f"{name}_{i}") 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 \ elif isinstance (layer, nn.LayerBase) or \
isinstance (layer, ModelBase): isinstance (layer, ModelBase):