From 1724f2afe7d088f9a5024440440d623fa637b758 Mon Sep 17 00:00:00 2001
From: Alex <defiler.off@gmail.com>
Date: Sun, 7 Mar 2021 13:09:35 +0300
Subject: [PATCH] add correct kimg for resume

---
 training/training_loop.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/training/training_loop.py b/training/training_loop.py
index 14836ad2e..d81b90582 100755
--- a/training/training_loop.py
+++ b/training/training_loop.py
@@ -245,7 +245,11 @@ def training_loop(
     if rank == 0:
         print(f'Training for {total_kimg} kimg...')
         print()
-    cur_nimg = 0
+    # Grab acual kmig from resume pkl
+    if (resume_pkl is not None) and (rank == 0):
+      cur_nimg = int(resume_pkl.split('-')[-1].split('.')[0])*1000
+    else: 
+      cur_nimg = 0
     cur_tick = 0
     tick_start_nimg = cur_nimg
     tick_start_time = time.time()