add date after transform in cashew (#168)

Signed-off-by: Carlos Gomes <[email protected]>
IBM · Sep 24, 2024 · 4640128 · 4640128
1 parent 7ae0125
commit 4640128
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/terratorch/datasets/m_cashew_plantation.py b/terratorch/datasets/m_cashew_plantation.py
@@ -82,9 +82,10 @@ def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
             date = self._get_date(h5file)
             mask = np.array(h5file["label"])
 
-        output = {"image": image.astype(np.float32), "date": date, "mask": mask}
+        output = {"image": image.astype(np.float32), "mask": mask}
 
         output = self.transform(**output)
+        output["date"] = date
         output["mask"] = output["mask"].long()
 
         return output