Fix stats checkpointing when no eval-during-train

Summary:
Bug introduced in D15977931.

Fixes #104.

Reviewed By: chandlerzuo

Differential Revision: D17498076

fbshipit-source-id: f9a8368c8e7217cd8ed1cd6d22f20e568de2d529

Author: lw

torchbiggraph/train.py

@@ -613,13 +613,16 @@ def swap_partitioned_embeddings(
         return io_bytes
 
     if rank == RANK_ZERO:
-        for stats in checkpoint_manager.maybe_read_stats():
-            yield (
-                stats["index"],
-                Stats.from_dict(stats["eval_stats_before"]),
-                Stats.from_dict(stats["stats"]),
-                Stats.from_dict(stats["eval_stats_after"]),
-            )
+        for stats_dict in checkpoint_manager.maybe_read_stats():
+            index: int = stats_dict["index"]
+            stats: Stats = Stats.from_dict(stats_dict["stats"])
+            eval_stats_before: Optional[Stats] = None
+            if "eval_stats_before" in stats_dict:
+                eval_stats_before = Stats.from_dict(stats_dict["eval_stats_before"])
+            eval_stats_after: Optional[Stats] = None
+            if "eval_stats_after" in stats_dict:
+                eval_stats_after = Stats.from_dict(stats_dict["eval_stats_after"])
+            yield (index, eval_stats_before, stats, eval_stats_after)
 
     # Start of the main training loop.
     for epoch_idx, edge_path_idx, edge_chunk_idx in iteration_manager:
@@ -771,14 +774,15 @@ def swap_partitioned_embeddings(
                 bucket_logger.info(f"Stats after training: {eval_stats_after}")
 
             # Add train/eval metrics to queue
-            checkpoint_manager.append_stats(
-                {
-                    "index": current_index,
-                    "eval_stats_before": eval_stats_before.to_dict(),
-                    "stats": stats.to_dict(),
-                    "eval_stats_after": eval_stats_after.to_dict(),
-                }
-            )
+            stats_dict = {
+                "index": current_index,
+                "stats": stats.to_dict(),
+            }
+            if eval_stats_before is not None:
+                stats_dict["eval_stats_before"] = eval_stats_before.to_dict()
+            if eval_stats_after is not None:
+                stats_dict["eval_stats_after"] = eval_stats_after.to_dict()
+            checkpoint_manager.append_stats(stats_dict)
             yield current_index, eval_stats_before, stats, eval_stats_after
 
         swap_partitioned_embeddings(cur_b, None)