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)