Remove references to depth_index in apply_hooks

Author: TylerYep

torchinfo/layer_info.py

@@ -29,7 +29,6 @@ def __init__(
         var_name: str,
         module: nn.Module,
         depth: int,
-        depth_index: int | None = None,
         parent_info: LayerInfo | None = None,
     ) -> None:
         # Identifying information
@@ -43,7 +42,7 @@ def __init__(
         # {layer name: {col_name: value_for_row}}
         self.inner_layers: dict[str, dict[ColumnSettings, Any]] = {}
         self.depth = depth
-        self.depth_index = depth_index
+        self.depth_index: int | None = None  # set at the very end
         self.executed = False
         self.parent_info = parent_info
         self.var_name = var_name

torchinfo/torchinfo.py

@@ -253,11 +253,12 @@ def forward_pass(
     summary_list: list[LayerInfo] = []
     hooks: list[RemovableHandle] | None = None if x is None else []
     named_module = (model_name, model)
-    apply_hooks(named_module, model, batch_dim, summary_list, {}, hooks, all_layers)
+    apply_hooks(named_module, model, batch_dim, summary_list, hooks, all_layers)
 
     if x is None:
         if not summary_list or summary_list[0].var_name != model_name:
             summary_list.insert(0, LayerInfo("", model, 0))
+        set_depth_index(summary_list)
         return summary_list
 
     kwargs = set_device(kwargs, device)
@@ -318,33 +319,14 @@ def add_missing_layers(
     if not set(b) - set(a):
         return
 
-    for tag, _, i2, j1, j2 in difflib.SequenceMatcher(None, a, b).get_opcodes():
+    for tag, _, _, j1, j2 in difflib.SequenceMatcher(None, a, b).get_opcodes():
         # Ignore all other layer types besides "insert".
         if tag == "insert":
-            depth_shifts: dict[int, int] = {}
             for i, info in enumerate(all_layers[j1:j2]):
-                # Find the correct depth_index for the new layer.
-                info.depth_index = 1
-                for prev_layer in reversed(summary_list[: j1 + i]):
-                    if (
-                        prev_layer.depth_index is not None
-                        and prev_layer.depth == info.depth
-                    ):
-                        info.depth_index = prev_layer.depth_index + 1
-                        break
-                if info.depth not in depth_shifts:
-                    depth_shifts[info.depth] = 0
-                depth_shifts[info.depth] += 1
-
                 info.calculate_num_params()
                 info.check_recursive(summary_list)
                 summary_list.insert(j1 + i, info)
 
-            # Shift depths forward for all existing later layers
-            for info in summary_list[i2 + (j2 - j1) :]:
-                if info.depth_index is not None:
-                    info.depth_index += depth_shifts.get(info.depth, 0)
-
 
 def validate_user_params(
     input_data: INPUT_DATA_TYPE | None,
@@ -491,7 +473,6 @@ def apply_hooks(
     orig_model: nn.Module,
     batch_dim: int | None,
     summary_list: list[LayerInfo],
-    idx: dict[int, int],
     hooks: list[RemovableHandle] | None,
     all_layers: list[LayerInfo],
     curr_depth: int = 0,
@@ -505,15 +486,14 @@ def apply_hooks(
     # Fallback is used if the layer's pre-hook is never called, for example in
     # ModuleLists or Sequentials.
     var_name, module = named_module
-    info = LayerInfo(var_name, module, curr_depth, None, parent_info)
+    info = LayerInfo(var_name, module, curr_depth, parent_info)
     all_layers.append(info)
 
     def pre_hook(module: nn.Module, inputs: Any) -> None:
         """Create a LayerInfo object to aggregate information about that layer."""
         del inputs
         nonlocal info
-        idx[curr_depth] = idx.get(curr_depth, 0) + 1
-        info = LayerInfo(var_name, module, curr_depth, idx[curr_depth], parent_info)
+        info = LayerInfo(var_name, module, curr_depth, parent_info)
         info.calculate_num_params()
         info.check_recursive(summary_list)
         summary_list.append(info)
@@ -543,7 +523,6 @@ def hook(module: nn.Module, inputs: Any, outputs: Any) -> None:
             orig_model,
             batch_dim,
             summary_list,
-            idx,
             hooks,
             all_layers,
             curr_depth + 1,