Apply suggestions from code review

Co-authored-by: Mihir Patel <mihir.v.patel7@gmail.com>
mosaicml · Sep 10, 2024 · 7c191d8 · 7c191d8
1 parent 991323c
commit 7c191d8
Showing 1 changed file with 2 additions and 0 deletions.
diff --git 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,