Fix epoch reporting when restoring checkpoint (#1075)

Summary: Pull Request resolved: fairinternal/fairseq-py#1075 Differential Revision: D20322672 Pulled By: myleott fbshipit-source-id: 8845a10b10442bed77670b7740afa271123a3b5d
facebookresearch · Mar 7, 2020 · 86916e0 · 86916e0
1 parent adb5b9c
commit 86916e0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/fairseq/checkpoint_utils.py b/fairseq/checkpoint_utils.py
@@ -352,7 +352,7 @@ def _upgrade_state_dict(state):
         state["args"].dataset_impl = "lazy"
     # epochs start at 1
     state["extra_state"]["train_iterator"]["epoch"] = max(
-        getattr(state["extra_state"]["train_iterator"], "epoch", 1),
+        state["extra_state"]["train_iterator"].get("epoch", 1),
         1,
     )