tmp: regime fix for pytorch-lightning 2.0

ZettaAI · Oct 18, 2023 · b563936 · b563936
1 parent 2effad0
commit b563936
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/zetta_utils/training/lightning/regimes/alignment/base_encoder.py b/zetta_utils/training/lightning/regimes/alignment/base_encoder.py
@@ -102,7 +102,7 @@ def log_results(mode: str, title_suffix: str = "", **kwargs):
     def validation_epoch_start(self, _):  # pylint: disable=no-self-use
         seed_everything(42)
 
-    def validation_epoch_end(self, _):
+    def on_validation_epoch_end(self):
         self.log_results(
             "val",
             "worst",