From 7c191d87ba601e8b4fea1bbee0e11643dfde56fc Mon Sep 17 00:00:00 2001
From: bigning <ning.wang@databricks.com>
Date: Tue, 10 Sep 2024 11:04:38 -0700
Subject: [PATCH] Apply suggestions from code review

Co-authored-by: Mihir Patel <mihir.v.patel7@gmail.com>
---
 composer/utils/checkpoint.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/composer/utils/checkpoint.py b/composer/utils/checkpoint.py
index 80f1d38296..c6f5af15ca 100644
--- a/composer/utils/checkpoint.py
+++ b/composer/utils/checkpoint.py
@@ -609,6 +609,7 @@ def dist_cp_load(
 ):
     if version.parse(torch.__version__) >= version.parse('2.4.0'):
         if version.parse(torch.__version__) < version.parse('2.4.1'):
+            # PyTorch 2.4.0
             from torch.distributed.checkpoint.utils import CheckpointException
             try:
                 dist_cp.load(
@@ -645,6 +646,7 @@ def dist_cp_load(
                 else:
                     raise e
         else:
+            # PyTorch 2.4.1
             dist_cp.load(
                 state_dict=state_dict,
                 storage_reader=storage_reader,