Merge branch 'f/Fix_augmentation_setup' into 'main'

Uncomment self.augmentations in image classifier See merge request es/ai/hannah/hannah!372
ekut-es · Feb 27, 2024 · a043851 · a043851
2 parents ebebd62 + cf7159c
commit a043851
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/hannah/modules/vision/base.py b/hannah/modules/vision/base.py
@@ -286,10 +286,11 @@ def setup_augmentations(self, pipeline_configs):
 
         self.default_augmentation = torch.nn.Sequential(*default_augment)
         augmentations = {k: torch.nn.Sequential(*v) for k, v in augmentations.items()}
+
+        self.augmentations = torch.nn.ModuleDict(augmentations)
 
         return augmentations
 
-        # self.augmentations = torch.nn.ModuleDict(augmentations)
 
     def _get_dataloader(self, dataset, unlabeled_data=None, shuffle=False):
         batch_size = self.hparams["batch_size"]