From 5b08bf068ec2a6ba407d4644e80743314693f950 Mon Sep 17 00:00:00 2001 From: jyz Date: Sun, 5 Jul 2026 21:15:08 +0000 Subject: [PATCH] fix(nn.utils.prune): handle empty PruningContainer --- test/nn/test_pruning.py | 20 ++++++++++++++++++++ torch/nn/utils/prune.py | 8 ++++++-- 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/test/nn/test_pruning.py b/test/nn/test_pruning.py index dbe288464efbf..9a2c3b28e88cb 100644 --- a/test/nn/test_pruning.py +++ b/test/nn/test_pruning.py @@ -392,6 +392,26 @@ def test_pruning_container(self): with self.assertRaises(TypeError): container.add_pruning_method("ugh") + def test_pruning_container_add_to_empty_adopts_tensor_name(self): + container = prune.PruningContainer() + + p = prune.L1Unstructured(amount=2) + p._tensor_name = "test" + container.add_pruning_method(p) + + self.assertEqual(container._tensor_name, "test") + self.assertEqual(len(container), 1) + + q = prune.RandomUnstructured(amount=2) + q._tensor_name = "test" + container.add_pruning_method(q) + self.assertEqual(len(container), 2) + + r = prune.L1Unstructured(amount=2) + r._tensor_name = "another_test" + with self.assertRaises(ValueError): + container.add_pruning_method(r) + def test_pruning_container_compute_mask(self): r"""Test `compute_mask` of pruning container with a known `t` and `default_mask`. Indirectly checks that Ln structured pruning is diff --git a/torch/nn/utils/prune.py b/torch/nn/utils/prune.py index 240fce1c32069..2d7b1385a14f7 100644 --- a/torch/nn/utils/prune.py +++ b/torch/nn/utils/prune.py @@ -300,10 +300,14 @@ def add_pruning_method(self, method) -> None: # check that we're adding a pruning method to the container if not isinstance(method, BasePruningMethod) and method is not None: raise TypeError(f"{type(method)} is not a BasePruningMethod subclass") - elif method is not None and self._tensor_name != method._tensor_name: + tensor_name = getattr(self, "_tensor_name", None) + if method is not None and tensor_name is None: + self._tensor_name = method._tensor_name + tensor_name = self._tensor_name + elif method is not None and tensor_name != method._tensor_name: raise ValueError( "Can only add pruning methods acting on " - f"the parameter named '{self._tensor_name}' to PruningContainer {self}." + f"the parameter named '{tensor_name}' to PruningContainer {self}." + f" Found '{method._tensor_name}'" ) # if all checks passed, add to _pruning_methods tuple