🎨 Consistently overwrite config values with calculated values when building Zebra models

Author: jemrobinson

ice_station_zebra/models/encode_process_decode.py

@@ -27,28 +27,22 @@ def __init__(
         # Add one encoder per dataset
         self.encoders = [
             hydra.utils.instantiate(
-                dict(
-                    {"input_space": input_space, "latent_space": latent_space_},
-                    **encoder,
-                ),
+                dict(**encoder)
+                | {"input_space": input_space, "latent_space": latent_space_}
             )
             for input_space in self.input_spaces
         ]
 
         # Add a processor
         self.processor = hydra.utils.instantiate(
-            dict(
-                {"n_latent_channels": latent_space_.channels * len(self.encoders)},
-                **processor,
-            ),
+            dict(**processor)
+            | {"n_latent_channels": latent_space_.channels * len(self.encoders)}
         )
 
         # Add a decoder
         self.decoder = hydra.utils.instantiate(
-            dict(
-                {"latent_space": latent_space_, "output_space": self.output_space},
-                **decoder,
-            ),
+            dict(**decoder)
+            | {"latent_space": latent_space_, "output_space": self.output_space}
         )
 
         # Register all modules that need to be trained

ice_station_zebra/models/zebra_model.py

@@ -46,10 +46,7 @@ def forward(self, inputs: LightningBatch) -> torch.Tensor:
     def configure_optimizers(self) -> Optimizer:
         """Construct the optimizer from the config"""
         return hydra.utils.instantiate(
-            dict(
-                {"params": self.model_list.parameters()},
-                **self.optimizer_cfg,
-            ),
+            dict(**self.optimizer_cfg) | {"params": self.model_list.parameters()}
         )
 
     def loss(self, output: torch.Tensor, target: torch.Tensor) -> torch.Tensor: