rename reference state dicts

pytorch · Jul 12, 2023 · 0998efa · 0998efa
1 parent 11b2c49
commit 0998efa
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/local_test_checkpoint.py b/test/local_test_checkpoint.py
@@ -96,8 +96,8 @@ def run_worker(args: List[str | int]) -> None:
 
     # Take an optimization step
     optimizer.step()
-    submod_ref = deepcopy(stage.submod.state_dict())
-    optim_ref = deepcopy(optimizer.state_dict())
+    ref_state_dict = deepcopy(stage.submod.state_dict())
+    ref_optim_state_dict = deepcopy(optimizer.state_dict())
     save_checkpoint(stage, CKPT_DIR, optimizer)
 
     # save index file in rank 0