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
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.