Fix the model summary table

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, module in enumerate(self.encoders):
+            self.add_module(f"encoder_{idx}", module)
 
         # 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,8 @@
+import itertools
 from abc import ABC, abstractmethod
 
 import hydra
 import torch
-import torch.nn as nn
 from lightning import LightningModule
 from omegaconf import DictConfig
 from torch.optim import Optimizer
@@ -28,9 +28,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,7 +43,12 @@ 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": itertools.chain(
+                    *[module.parameters() for module in self.children()]
+                )
+            }
         )
 
     def loss(self, output: torch.Tensor, target: torch.Tensor) -> torch.Tensor: