remove outputs from cudacallback

2022-09-07 03:51:04 -07:00 · 2022-09-07 03:51:04 -07:00 · d283e04b8b
parent f7bb9b7bbc
commit d283e04b8b
1 changed files with 1 additions and 1 deletions
--- a/main.py
+++ b/main.py
@ -400,7 +400,7 @@ class CUDACallback(Callback):
        torch.cuda.synchronize(trainer.root_gpu)
        self.start_time = time.time()

-    def on_train_epoch_end(self, trainer, pl_module, outputs):
+    def on_train_epoch_end(self, trainer, pl_module):
        torch.cuda.synchronize(trainer.root_gpu)
        max_memory = torch.cuda.max_memory_allocated(trainer.root_gpu) / 2 ** 20
        epoch_time = time.time() - self.start_time