Removing hard coded name

mosaicml · Jan 6, 2024 · cf6d9e9 · cf6d9e9
1 parent d5ad317
commit cf6d9e9
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/composer/utils/checkpoint.py b/composer/utils/checkpoint.py
@@ -497,7 +497,8 @@ def read_data(self, plan: LoadPlan, planner: LoadPlanner):
                 log.debug('Strip _pgidx from optimizer state dict keys')
                 local_idx = f'_pgidx{dist.get_local_rank()}'
                 log.debug('Get ptr to optimizer state dict')
-                optim_state_dict = optim_state['optimizers']['DecoupledLionW']['state']
+                optim_name = list(optim_state['optimizers'].keys())[0]
+                optim_state_dict = optim_state['optimizers'][optim_name]['state']
                 log.debug('Loop over optimizer state dict keys')
                 for key in list(optim_state_dict.keys()):
                     log.debug(f'Stripping {local_idx} from {key=}')