Fix callback typing.

Lightning-AI · Aug 7, 2020 · f9f762e · f9f762e
1 parent c93d8c8
commit f9f762e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/trainer/test_states.py b/tests/trainer/test_states.py
@@ -18,7 +18,7 @@ def on_batch_start(self, trainer, pl_module):
         if self.snapshot_method == 'on_batch_start':
             self.trainer_state = trainer.state
 
-    def on_test_batch_start(self, trainer, pl_module):
+    def on_test_batch_start(self, trainer, pl_module, batch, batch_idx, dataloader_idx):
         if self.snapshot_method == 'on_test_batch_start':
             self.trainer_state = trainer.state