Fixing validation step for PL 1

rob-bru · Jun 30, 2023 · 061d11d · 061d11d
1 parent 2796c81
commit 061d11d
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/main.py b/main.py
@@ -469,9 +469,8 @@ def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
             self.log_img(pl_module, batch, batch_idx, split="train")
 
     @rank_zero_only
-    # def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
     def on_validation_batch_end(
-        self, trainer, pl_module, outputs, batch, batch_idx, **kwargs
+        self, trainer, pl_module, outputs, batch, batch_idx, *args, **kwargs
     ):
         if not self.disabled and pl_module.global_step > 0:
             self.log_img(pl_module, batch, batch_idx, split="val")