From a4f60690292987679b9dc2a76108539188dbd0df Mon Sep 17 00:00:00 2001
From: Nickolay Shmyrev <nshmyrev@gmail.com>
Date: Mon, 3 Jul 2023 18:20:16 +0200
Subject: [PATCH] Also set epoch

---
 icefall/rnn_lm/train.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/icefall/rnn_lm/train.py b/icefall/rnn_lm/train.py
index 6ff4f16ec..80f8c238f 100755
--- a/icefall/rnn_lm/train.py
+++ b/icefall/rnn_lm/train.py
@@ -278,7 +278,7 @@ def load_checkpoint_if_available(
     elif params.start_epoch > 1:
         filename = params.exp_dir / f"epoch-{params.start_epoch-1}.pt"
     else:
-        return
+        return None
 
     logging.info(f"Loading checkpoint: {filename}")
     saved_params = load_checkpoint(
@@ -298,6 +298,10 @@ def load_checkpoint_if_available(
     for k in keys:
         params[k] = saved_params[k]
 
+    if params.start_batch > 0:
+        if "cur_epoch" in saved_params:
+            params["start_epoch"] = saved_params["cur_epoch"]
+
     return saved_params