make get_layers_from_model_by_type in torch more readable.

sony · Sep 25, 2024 · 0a228f5 · 0a228f5
1 parent 16a47bf
commit 0a228f5
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/tests/pytorch_tests/utils.py b/tests/pytorch_tests/utils.py
@@ -39,10 +39,11 @@ def get_layers_from_model_by_type(model: torch.nn.Module,
         List of layers of type layer_type from the model.
     """
     match_layer_type = lambda _layer: layer_type in [type(_layer), _layer]
+    matched_list = [layer[1] for layer in model.named_children() if match_layer_type(layer[1])]
     if include_wrapped_layers:
-        return [layer[1] for layer in model.named_children() if match_layer_type(layer[1]) or
-                (isinstance(layer[1], PytorchQuantizationWrapper) and match_layer_type(layer[1].layer))]
-    return [layer[1] for layer in model.named_children() if match_layer_type(layer[1])]
+        matched_list.extend([layer[1] for layer in model.named_children()
+                             if (isinstance(layer[1], PytorchQuantizationWrapper) and match_layer_type(layer[1].layer))])
+    return matched_list
 
 
 def count_model_prunable_params(model: torch.nn.Module) -> int: