Merge branch 'master' of github.com:amarczew/pytorch_model_summary in…

…to master
amarczew · Aug 30, 2020 · 9ff3877 · 9ff3877
2 parents 427a800 + 1742743
commit 9ff3877
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -1,4 +1,7 @@
-## Pytorch Model Summary
+## Pytorch Model Summary -- Keras style `model.summary()` for PyTorch
+[![PyPI version fury.io](https://badge.fury.io/py/pytorch-model-summary.svg)](https://pypi.python.org/pypi/pytorch-model-summary/)
+[![Downloads](https://pepy.tech/badge/pytorch-model-summary)](https://pepy.tech/project/pytorch-model-summary)
+
 
 It is a Keras style model.summary() implementation for PyTorch
 

diff --git a/pytorch_model_summary/hierarchical_summary.py b/pytorch_model_summary/hierarchical_summary.py
@@ -19,6 +19,8 @@ def repr(model):
         child_lines = []
         total_params = 0
         for key, module in model._modules.items():
+            if module is None:
+                continue
             mod_str, num_params = repr(module)
             mod_str = _addindent(mod_str, 2)
             child_lines.append('(' + key + '): ' + mod_str)