Convert get_param_count to @staticmethod

Author: TylerYep

torchinfo/layer_info.py

@@ -126,17 +126,10 @@ def nested_list_size(inputs: Sequence[Any]) -> tuple[list[int], int]:
 
         return size, elem_bytes
 
-    def get_layer_name(self, show_var_name: bool, show_depth: bool) -> str:
-        layer_name = self.class_name
-        if show_var_name and self.var_name:
-            layer_name += f" ({self.var_name})"
-        if show_depth and self.depth > 0:
-            layer_name += f": {self.depth}"
-            if self.depth_index is not None:
-                layer_name += f"-{self.depth_index}"
-        return layer_name
-
-    def get_param_count(self, name: str, param: torch.Tensor) -> tuple[int, str]:
+    @staticmethod
+    def get_param_count(
+        module: nn.Module, name: str, param: torch.Tensor
+    ) -> tuple[int, str]:
         """
         Get count of number of params, accounting for mask.
 
@@ -146,12 +139,22 @@ def get_param_count(self, name: str, param: torch.Tensor) -> tuple[int, str]:
         """
         if name.endswith("_orig"):
             without_suffix = name[:-5]
-            pruned_weights = rgetattr(self.module, f"{without_suffix}_mask")
+            pruned_weights = rgetattr(module, f"{without_suffix}_mask")
             if pruned_weights is not None:
                 parameter_count = int(torch.sum(pruned_weights))
                 return parameter_count, without_suffix
         return param.nelement(), name
 
+    def get_layer_name(self, show_var_name: bool, show_depth: bool) -> str:
+        layer_name = self.class_name
+        if show_var_name and self.var_name:
+            layer_name += f" ({self.var_name})"
+        if show_depth and self.depth > 0:
+            layer_name += f": {self.depth}"
+            if self.depth_index is not None:
+                layer_name += f"-{self.depth_index}"
+        return layer_name
+
     def calculate_num_params(self) -> None:
         """
         Set num_params, trainable, inner_layers, and kernel_size
@@ -161,7 +164,7 @@ def calculate_num_params(self) -> None:
         for name, param in self.module.named_parameters():
             if is_lazy(param):
                 continue
-            cur_params, name = self.get_param_count(name, param)
+            cur_params, name = self.get_param_count(self.module, name, param)
             self.param_bytes += param.element_size() * cur_params
 
             self.num_params += cur_params
@@ -194,7 +197,7 @@ def calculate_macs(self) -> None:
         i.e., taking the batch-dimension into account.
         """
         for name, param in self.module.named_parameters():
-            cur_params, name = self.get_param_count(name, param)
+            cur_params, name = self.get_param_count(self.module, name, param)
             if name in ("weight", "bias"):
                 # ignore C when calculating Mult-Adds in ConvNd
                 if "Conv" in self.class_name: