From 1431436ecf0078c67c3f17ce039c3fa7c671a915 Mon Sep 17 00:00:00 2001
From: Yuru Jia <91590963+yurujaja@users.noreply.github.com>
Date: Fri, 13 Sep 2024 10:28:45 +0200
Subject: [PATCH] minor fix

---
 engine/data_preprocessor.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/engine/data_preprocessor.py b/engine/data_preprocessor.py
index b85fe0bd..7eacfc8d 100644
--- 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)
\ No newline at end of file
+        super().__init__(dataset, cfg, local_cfg)