From 0c1a16ea5b1452e4c81ce8ed17613f1761f4db21 Mon Sep 17 00:00:00 2001
From: Frank Lee <somerlee.9@gmail.com>
Date: Thu, 28 Jul 2022 09:29:30 +0800
Subject: [PATCH] [util] standard checkpoint function naming (#1377)

---
 colossalai/utils/checkpoint/module_checkpoint.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/colossalai/utils/checkpoint/module_checkpoint.py b/colossalai/utils/checkpoint/module_checkpoint.py
index 43c518394..dc01b3b54 100644
--- a/colossalai/utils/checkpoint/module_checkpoint.py
+++ b/colossalai/utils/checkpoint/module_checkpoint.py
@@ -6,7 +6,7 @@ from colossalai.utils.checkpoint.utils import gather_tensor, scatter_tensor
 from typing import Optional
 
 
-def save_checkpoint(dire: str,
+def save_checkpoint(path: str,
                     epoch: int,
                     model: torch.nn.Module,
                     optimizer: Optional[ColossalaiOptimizer] = None,
@@ -16,7 +16,7 @@ def save_checkpoint(dire: str,
     """save_checkpoint 
     save a model, whose parameters are `ColoTensor`s.
     Args:
-        dire (str): directory to save the checkpoint files.
+        path (str): directory to save the checkpoint files.
         epoch (int): the number of epoch
         model (torch.nn.Module): a torch module initialized by ColoInitContext
         optimizer (ColossalaiOptimizer, optional): optimizers. Defaults to None.
@@ -39,7 +39,7 @@ def save_checkpoint(dire: str,
                 delattr(v, 'save_ready')
         # model saving
         save_state = {'epoch': epoch, 'model': model_state}
-        torch.save(save_state, dire + '/epoch_{}_model.pth'.format(epoch), *args, **kwargs)
+        torch.save(save_state, path + '/epoch_{}_model.pth'.format(epoch), *args, **kwargs)
 
     # delete old dicts
     del model_state
@@ -57,7 +57,7 @@ def save_checkpoint(dire: str,
 
         if rank == 0:
             save_state = {'epoch': epoch, 'optim': optim_state}
-            torch.save(save_state, dire + '/epoch_{}_optim.pth'.format(epoch), *args, **kwargs)
+            torch.save(save_state, path + '/epoch_{}_optim.pth'.format(epoch), *args, **kwargs)
             # recover colo tensors in rank0
             for k, v in optimizer.state_dict()['state'].items():
                 for n, t in v.items():
@@ -71,7 +71,7 @@ def save_checkpoint(dire: str,
         dist.barrier()
 
 
-def load_checkpoint(dire,
+def load_checkpoint(path,
                     epoch: int,
                     model: torch.nn.Module,
                     optimizer: Optional[ColossalaiOptimizer] = None,
@@ -81,7 +81,7 @@ def load_checkpoint(dire,
     """load_checkpoint 
     load a model, whose parameters are `ColoTensor`s.
     Args:
-        dire (_type_): _description_
+        path (_type_): _description_
         epoch (int): _description_
         rank (int): _description_
         model (torch.nn.Module): _description_
@@ -96,7 +96,7 @@ def load_checkpoint(dire,
             gather_tensor(p)
 
     if rank == 0:
-        load_state = torch.load(dire + '/epoch_{}_model.pth'.format(epoch), *args, **kwargs)
+        load_state = torch.load(path + '/epoch_{}_model.pth'.format(epoch), *args, **kwargs)
         model.load_state_dict(load_state['model'])
     dist.barrier()
 
@@ -118,7 +118,7 @@ def load_checkpoint(dire,
                     gather_tensor(t)
 
         if rank == 0:
-            colo_checkpoint = torch.load(dire + '/epoch_{}_optim.pth'.format(epoch), *args, **kwargs)
+            colo_checkpoint = torch.load(path + '/epoch_{}_optim.pth'.format(epoch), *args, **kwargs)
             optimizer.load_state_dict(colo_checkpoint['optim'])
         dist.barrier()