Some fixes for compatibility with PL master

Lightning-Universe · Jan 7, 2022 · eb47dc5 · eb47dc5
1 parent 3b2d868
commit eb47dc5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/flash/core/data/io/input_transform.py b/flash/core/data/io/input_transform.py
@@ -921,7 +921,7 @@ def _resolve_transforms(self, running_stage: RunningStage) -> Optional[Dict[str,
         from flash.core.data.data_pipeline import DataPipeline
 
         transforms_out = {}
-        stage = _STAGES_PREFIX[running_stage.value]
+        stage = _STAGES_PREFIX[running_stage]
 
         # iterate over all transforms hook name
         for transform_name in InputTransformPlacement:
@@ -997,7 +997,7 @@ def _check_transforms(
         if transform is None:
             return transform
 
-        keys_diff = set(transform.keys()).difference([v for v in InputTransformPlacement])
+        keys_diff = set(transform.keys()).difference([v.value for v in InputTransformPlacement])
 
         if len(keys_diff) > 0:
             raise MisconfigurationException(