dont wait until 2nd epoch to measure perf

a9d3dc9f · Marta · f1a9d2b4 · a9d3dc9f
Commit a9d3dc9f authored Nov 02, 2021 by Marta
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 4 deletions

openfold/utils/logger.py openfold/utils/logger.py +2 -4

No files found.
--- a/openfold/utils/logger.py
+++ b/openfold/utils/logger.py
@@ -43,12 +43,10 @@ class PerformanceLoggingCallback(Callback):
            self.timestamps.append(time.time())

    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx, dataloader_idx):
-        if trainer.current_epoch == 1:
-            self.do_step()
+        self.do_step()

    def on_test_batch_start(self, trainer, pl_module, batch, batch_idx, dataloader_idx):
-        if trainer.current_epoch == 1:
-            self.do_step()
+        self.do_step()

    def process_performance_stats(self, deltas):
        def _round3(val):