fix logs

abcucberkeley · Nov 12, 2024 · 2420b70 · 2420b70
1 parent 57956fd
commit 2420b70
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/train.py b/src/train.py
@@ -422,7 +422,7 @@ def trace_handler(prof):
                 cuda_memory_allocated = np.mean(step_vram)
                 max_cuda_memory_allocated = torch.cuda.max_memory_allocated() / (1024 ** 3)
 
-                logger.info(f"│ training_epoch:           \t {epoch}/{remaining_epochs}")
+                logger.info(f"│ training_epoch:           \t {epoch}/{config['epochs']}")
                 logger.info(f"│ epoch_loss:               \t {loss:.4g}")
                 logger.info(f"│ epoch_mse:                \t {mse:.4g}")
                 logger.info(f"│ epoch_lr:                 \t {lr:.4g}")