Fix hidden nodes' remove in Retiarii (#3736)

nni/retiarii/nn/pytorch/mutator.py

@@ -33,7 +33,7 @@ class LayerChoiceMutator(Mutator):
             model.get_node_by_name(node.name).update_operation(Cell(node.operation.cell_name))
 
             # remove redundant nodes
-            for rm_node in target.hidden_nodes:
+            for rm_node in list(target.hidden_nodes):  # remove from a list on the fly will cause issues
                 if rm_node.name != chosen_node.name:
                     rm_node.remove()
 

test/ut/retiarii/test_highlevel_apis.py

@@ -93,6 +93,24 @@ class GraphIR(unittest.TestCase):
         self.assertEqual(self._get_converted_pytorch_model(model2)(torch.randn(1, 3, 3, 3)).size(),
                          torch.Size([1, 5, 3, 3]))
 
+    def test_layer_choice_multiple(self):
+        @self.get_serializer()
+        class Net(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.module = nn.LayerChoice([nn.Conv2d(3, i, kernel_size=1) for i in range(1, 11)])
+
+            def forward(self, x):
+                return self.module(x)
+
+        model, mutators = self._get_model_with_mutators(Net())
+        self.assertEqual(len(mutators), 1)
+        mutator = mutators[0].bind_sampler(EnumerateSampler())
+        for i in range(1, 11):
+            model_new = mutator.apply(model)
+            self.assertEqual(self._get_converted_pytorch_model(model_new)(torch.randn(1, 3, 3, 3)).size(),
+                             torch.Size([1, i, 3, 3]))
+
     def test_input_choice(self):
         @self.get_serializer()
         class Net(nn.Module):