hotfix (#261)

* hotfix * update
Lightning-Universe · May 5, 2021 · b3c049c · b3c049c
1 parent ec404e8
commit b3c049c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/flash/data/data_pipeline.py b/flash/data/data_pipeline.py
@@ -575,7 +575,12 @@ def __init__(self, func_to_wrap: Callable, model: 'Task') -> None:
     def __call__(self, *args, **kwargs):
         outputs = self.func(*args, **kwargs)
 
-        internal_running_state = self.internal_mapping[self.model.trainer._running_stage]
+        try:
+            stage = self.model.trainer._running_stage
+        except AttributeError:
+            stage = self.model.trainer.state.stage
+
+        internal_running_state = self.internal_mapping[stage]
         additional_func = self._stage_mapping.get(internal_running_state, None)
 
         if additional_func:

diff --git a/flash/video/classification/model.py b/flash/video/classification/model.py
@@ -43,6 +43,7 @@
 class VideoClassifierFinetuning(BaseFinetuning):
 
     def __init__(self, num_layers: int = 5, train_bn: bool = True, unfreeze_epoch: int = 1):
+        super().__init__()
         self.num_layers = num_layers
         self.train_bn = train_bn
         self.unfreeze_epoch = unfreeze_epoch