Set depth_index at the very end

Author: TylerYep

tests/test_output/module_dict.out

@@ -5,9 +5,9 @@ ModuleDictModel                          --                        --
 ├─ModuleDict: 1-1                        --                        --
 │    └─Conv2d: 2-1                       [1, 10, 1, 1]             910
 │    └─MaxPool2d: 2-2                    --                        --
-├─ModuleDict: 1-3                        --                        --
-│    └─LeakyReLU: 2-4                    [1, 10, 1, 1]             --
-│    └─PReLU: 2-5                        --                        1
+├─ModuleDict: 1-2                        --                        --
+│    └─LeakyReLU: 2-3                    [1, 10, 1, 1]             --
+│    └─PReLU: 2-4                        --                        1
 ==========================================================================================
 Total params: 911
 Trainable params: 911

torchinfo/torchinfo.py

@@ -289,11 +289,19 @@ def forward_pass(
         summary_list.insert(0, LayerInfo("", model, 0))
 
     add_missing_layers(summary_list, all_layers)
+    set_depth_index(summary_list)
 
     _cached_forward_pass[model_name] = summary_list
     return summary_list
 
 
+def set_depth_index(summary_list: list[LayerInfo]) -> None:
+    idx: dict[int, int] = {}
+    for layer in summary_list:
+        idx[layer.depth] = idx.get(layer.depth, 0) + 1
+        layer.depth_index = idx[layer.depth]
+
+
 def add_missing_layers(
     summary_list: list[LayerInfo], all_layers: list[LayerInfo]
 ) -> None: