From b46b1e70c6dcf0c40bdd9ec909c4ca6907e2ce12 Mon Sep 17 00:00:00 2001 From: Christian Steinmeyer Date: Fri, 18 Aug 2023 14:09:28 +0200 Subject: [PATCH] remove early out If there are multiple input layers (and therefore paths), this early out might return an empty set before all paths are checked. By simply removing it, the path that doesn't have a next layer ends, nothing is added to the `found_layers` and the other paths continue to be checked. --- .../python/core/sparsity/keras/pruning_policy.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tensorflow_model_optimization/python/core/sparsity/keras/pruning_policy.py b/tensorflow_model_optimization/python/core/sparsity/keras/pruning_policy.py index 39c5fe1fe..197253d92 100644 --- a/tensorflow_model_optimization/python/core/sparsity/keras/pruning_policy.py +++ b/tensorflow_model_optimization/python/core/sparsity/keras/pruning_policy.py @@ -124,8 +124,6 @@ def _lookup_layers(self, source_layers, stop_fn, next_fn): found_layers.add(layer) else: next_layers = next_fn(layer) - if not next_layers: - return set() for next_layer in next_layers: if next_layer not in used_layers: used_layers.add(next_layer)