Fix silly type check and improve readability of for loop

Author: stefanradev93

bayesflow/datasets/disk_dataset.py

@@ -92,9 +92,11 @@ def __getitem__(self, item) -> dict[str, np.ndarray]:
 
         batch = tree_stack(batch)
 
-        if isinstance(self.augmentations, Mapping):
-            for key in self.augmentations:
-                batch[key] = self.augmentations[key](batch[key])
+        if self.augmentations is None:
+            pass
+        elif isinstance(self.augmentations, Mapping):
+            for key, fn in self.augmentations.items():
+                batch[key] = fn(batch[key])
         elif isinstance(self.augmentations, Callable):
             batch = self.augmentations(batch)
         else:

bayesflow/datasets/offline_dataset.py

@@ -102,9 +102,11 @@ def __getitem__(self, item: int) -> dict[str, np.ndarray]:
             for key, value in self.data.items()
         }
 
-        if isinstance(self.augmentations, Mapping):
-            for key in self.augmentations:
-                batch[key] = self.augmentations[key](batch[key])
+        if self.augmentations is None:
+            pass
+        elif isinstance(self.augmentations, Mapping):
+            for key, fn in self.augmentations.items():
+                batch[key] = fn(batch[key])
         elif isinstance(self.augmentations, Callable):
             batch = self.augmentations(batch)
         else:

bayesflow/datasets/online_dataset.py

@@ -76,9 +76,11 @@ def __getitem__(self, item: int) -> dict[str, np.ndarray]:
         """
         batch = self.simulator.sample((self.batch_size,))
 
-        if isinstance(self.augmentations, Mapping):
-            for key in self.augmentations:
-                batch[key] = self.augmentations[key](batch[key])
+        if self.augmentations is None:
+            pass
+        elif isinstance(self.augmentations, Mapping):
+            for key, fn in self.augmentations.items():
+                batch[key] = fn(batch[key])
         elif isinstance(self.augmentations, Callable):
             batch = self.augmentations(batch)
         else: