Skip to content

Commit

Permalink
Do not show "needs to register block" warning for registered blocks. (a…
Browse files Browse the repository at this point in the history
  • Loading branch information
aalexandrov authored and szha committed Aug 12, 2018
1 parent cc132ef commit d710c4e
Showing 1 changed file with 7 additions and 6 deletions.
13 changes: 7 additions & 6 deletions python/mxnet/gluon/block.py
Original file line number Diff line number Diff line change
Expand Up @@ -208,26 +208,27 @@ def __setattr__(self, name, value):
super(Block, self).__setattr__(name, value)

def _check_container_with_block(self):
def _find_block_in_container(data):
children = set(self._children.values())
def _find_unregistered_block_in_container(data):
# Find whether a nested container structure contains Blocks
if isinstance(data, (list, tuple)):
for ele in data:
if _find_block_in_container(ele):
if _find_unregistered_block_in_container(ele):
return True
return False
elif isinstance(data, dict):
for _, v in data.items():
if _find_block_in_container(v):
if _find_unregistered_block_in_container(v):
return True
return False
elif isinstance(data, Block):
return True
return not data in children
else:
return False
for k, v in self.__dict__.items():
if isinstance(v, (list, tuple, dict)) and not (k.startswith('__') or k == '_children'):
if _find_block_in_container(v):
warnings.warn('"{name}" is a container with Blocks. '
if _find_unregistered_block_in_container(v):
warnings.warn('"{name}" is an unregistered container with Blocks. '
'Note that Blocks inside the list, tuple or dict will not be '
'registered automatically. Make sure to register them using '
'register_child() or switching to '
Expand Down

0 comments on commit d710c4e

Please sign in to comment.