augmentation=none fix

InnopolisUni · Dec 7, 2023 · aa30bee · aa30bee
1 parent 45f67f7
commit aa30bee
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/innofw/utils/framework.py b/innofw/utils/framework.py
@@ -143,7 +143,7 @@ def is_albu(aug):
         return isinstance(aug, A.Compose) or isinstance(
             aug, A.core.transforms_interface.BasicTransform
         )
-    cfg = cfg['augmentations']
+    cfg = cfg.get('augmentations', cfg) if cfg is not None else cfg
     for key in keys:
         try:
             for v in cfg[key].values():