fix(nn.utils.prune): don't crash adding a method to an empty PruningContainer

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

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