Do not error if there are None values in ModuleLists

TylerYep · May 16, 2022 · e7be90b · e7be90b
1 parent 7def795
commit e7be90b
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 13 deletions.
diff --git a/tests/fixtures/models.py b/tests/fixtures/models.py
@@ -407,11 +407,14 @@ def __init__(self) -> None:
         self._layers.append(nn.Linear(5, 5))
         self._layers.append(ContainerChildModule())
         self._layers.append(nn.Linear(5, 5))
+        # Add None, but filter out this value later.
+        self._layers.append(None)  # type: ignore[arg-type]
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         out = x
         for layer in self._layers:
-            out = layer(out)
+            if layer is not None:
+                out = layer(out)
         return out
 
 

diff --git a/torchinfo/torchinfo.py b/torchinfo/torchinfo.py
@@ -566,18 +566,18 @@ def hook(module: nn.Module, inputs: Any, outputs: Any) -> None:
 
     # module.named_modules(remove_duplicate=False) doesn't work (infinite recursion).
     for name, mod in module._modules.items():  # pylint: disable=protected-access
-        assert mod is not None
-        child = (name, mod)
-        apply_hooks(
-            child,
-            orig_model,
-            batch_dim,
-            summary_list,
-            hooks,
-            all_layers,
-            curr_depth + 1,
-            info,
-        )
+        if mod is not None:
+            child = (name, mod)
+            apply_hooks(
+                child,
+                orig_model,
+                batch_dim,
+                summary_list,
+                hooks,
+                all_layers,
+                curr_depth + 1,
+                info,
+            )
 
 
 def clear_cached_forward_pass() -> None: