Class MyModel

python/tests/test_nn.py:324–328

        # Using leaf_modules in the update should always work
        class MyModel(nn.Module):
            def __init__(self):
                super().__init__()
                self.stuff = [nn.Linear(2, 2), 0, nn.Linear(2, 2)]
                self.more_stuff = {"hi": nn.Linear(2, 2), "bye": 0}

        m = MyModel()
        m.update_modules(m.leaf_modules())