minor fix

VMarsocci · Sep 13, 2024 · 1431436 · 1431436
1 parent 35e600a
commit 1431436
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/engine/data_preprocessor.py b/engine/data_preprocessor.py
@@ -17,7 +17,7 @@
 
 
 def get_collate_fn(cfg: omegaconf.DictConfig) -> Callable:
-    modalities = cfg.dataset.bands.keys()
+    modalities = cfg.encoder.input_bands.keys()
 
     def collate_fn(
         batch: dict[dict[str, torch.Tensor]]
@@ -602,4 +602,4 @@ def __init__(self, dataset, cfg, local_cfg):
         if not local_cfg:
             local_cfg = omegaconf.OmegaConf.create()
         local_cfg.size = cfg.encoder.input_size
-        super().__init__(dataset, cfg, local_cfg)
+        super().__init__(dataset, cfg, local_cfg)