🐛 Use the implicit self._modules instead of a manual ModuleList to keep track of child modules. This fixes the incorrectly displayed summary table

Author: jemrobinson

ice_station_zebra/models/encode_process_decode.py

@@ -32,6 +32,10 @@ def __init__(
             )
             for input_space in self.input_spaces
         ]
+        # We have to explicitly register each encoder as list[Module] will not be
+        # automatically picked up by PyTorch
+        for idx, encoder in enumerate(self.encoders):
+            self.add_module(f"encoder_{idx}", encoder)
 
         # Add a processor
         self.processor = hydra.utils.instantiate(
@@ -45,11 +49,6 @@ def __init__(
             | {"latent_space": latent_space_, "output_space": self.output_space}
         )
 
-        # Register all modules that need to be trained
-        self.model_list.extend(self.encoders)
-        self.model_list.append(self.processor)
-        self.model_list.append(self.decoder)
-
     def forward(self, inputs: LightningBatch) -> torch.Tensor:
         """Forward step of the model
 

ice_station_zebra/models/zebra_model.py

@@ -1,8 +1,9 @@
+import itertools
 from abc import ABC, abstractmethod
+from collections.abc import Iterator
 
 import hydra
 import torch
-import torch.nn as nn
 from lightning import LightningModule
 from omegaconf import DictConfig
 from torch.optim import Optimizer
@@ -28,9 +29,6 @@ def __init__(
         self.input_spaces = [DataSpace.from_dict(space) for space in input_spaces]
         self.output_space = DataSpace.from_dict(output_space)
 
-        # Initialise an empty module list
-        self.model_list = nn.ModuleList()
-
         # Store the optimizer config
         self.optimizer_cfg = optimizer
 
@@ -46,12 +44,17 @@ def forward(self, inputs: LightningBatch) -> torch.Tensor:
     def configure_optimizers(self) -> Optimizer:
         """Construct the optimizer from the config"""
         return hydra.utils.instantiate(
-            dict(**self.optimizer_cfg) | {"params": self.model_list.parameters()}
+            dict(**self.optimizer_cfg) | {"params": self.parameters()}
         )
 
     def loss(self, output: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         return torch.nn.functional.l1_loss(output, target)
 
+    def parameters(self) -> Iterator[torch.nn.Parameter]:
+        return itertools.chain(
+            *[module.parameters() for module in self._modules.values()]
+        )
+
     def test_step(
         self, batch: LightningBatch, batch_idx: int
     ) -> dict[str, torch.Tensor]: