From a2f57e6bc42ab5505aff7fafd0aea8728eac8ccb Mon Sep 17 00:00:00 2001
From: Bar <elhararb@gmail.com>
Date: Fri, 30 Nov 2018 02:28:35 +0200
Subject: [PATCH] Fix assign_layer_fq_names (#88)
Add support to models that contain named empty layers (i.e. have type: NoneType).
This fix catches the AttributeError that is raised when a named NoneType layer is detected, and the layer is ignored.
---
distiller/utils.py | 8 ++++++--
1 file changed, 6 insertions(+), 2 deletions(-)
diff --git a/distiller/utils.py b/distiller/utils.py
index 29e8bde..17d1097 100755
--- a/distiller/utils.py
+++ b/distiller/utils.py
@@ -72,8 +72,12 @@ def assign_layer_fq_names(container, name=None):
"""
is_leaf = True
for key, module in container._modules.items():
- is_leaf = False
- assign_layer_fq_names(module, ".".join([name, key]) if name is not None else key)
+ try:
+ assign_layer_fq_names(module, ".".join([name, key]) if name is not None else key)
+ is_leaf = False
+ except AttributeError:
+ if module is not None:
+ raise
if is_leaf:
container.distiller_name = name
--
GitLab