bandaid to cached dataloader

viscy/data/hcs_ram.py

@@ -117,7 +117,8 @@ def __getitem__(self, index: int) -> Sample:
         if norm_meta is not None:
             sample_images["norm_meta"] = norm_meta
         if self.transform:
-            sample_images = self.transform(sample_images)
+            # FIX ME: check why the transforms return a list?
+            sample_images = self.transform(sample_images)[0]
         if "weight" in sample_images:
             del sample_images["weight"]
         sample = {
@@ -185,6 +186,11 @@ def setup(self, stage: Literal["fit", "validate", "test", "predict"]) -> None:
             raise NotImplementedError(f"Stage {stage} is not supported")
 
     def _train_transform(self) -> list[Callable]:
+        """ Set the train augmentations
+
+        
+        """
+
         if self.augmentations:
             for aug in self.augmentations:
                 if isinstance(aug, MultiSampleTrait):
@@ -197,6 +203,10 @@ def _train_transform(self) -> list[Callable]:
                             f"transform type {type(aug)}."
                         )
                     self.train_patches_per_stack = num_samples
+        else:
+            self.augmentations=[]
+
+        _logger.info(f'Training augmentations: {self.augmentations}')
         return list(self.augmentations)
 
     def _fit_transform(self) -> tuple[Compose, Compose]: