Using variable number of modules as children when subclassing torch.nn.Module

PyTorch will register the submodules as children, if you use an nn.ModuleList instead of a Python list.
Use self.test_classes = nn.ModuleList() and it should work.

1 Like