Fixes

Lightning-Universe · Jul 9, 2021 · 40b2404 · 40b2404
1 parent fcffb56
commit 40b2404
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/flash/video/classification/model.py b/flash/video/classification/model.py
@@ -146,8 +146,8 @@ def on_train_epoch_start(self) -> None:
             encoded_dataset._video_sampler.set_epoch(self.trainer.current_epoch)
         super().on_train_epoch_start()
 
-    def step(self, batch: Any, batch_idx: int) -> Any:
-        return super().step((batch["video"], batch["label"]), batch_idx)
+    def step(self, batch: Any, batch_idx: int, metrics) -> Any:
+        return super().step((batch["video"], batch["label"]), batch_idx, metrics)
 
     def forward(self, x: Any) -> Any:
         x = self.backbone(x)