diff --git a/ppcls/optimizer/__init__.py b/ppcls/optimizer/__init__.py
index 9ec3bc897ecf66373a30f6f33cd81944d7ceb682..cc0041137b0abe436a2c4a137cd209cfe211236e 100644
--- a/ppcls/optimizer/__init__.py
+++ b/ppcls/optimizer/__init__.py
@@ -61,7 +61,7 @@ def build_optimizer(config, epochs, step_each_epoch, model_list=None):
     for optim_item in optim_config:
         # optim_cfg = {optim_name1: {scope: xxx, **optim_cfg}}
         # step1 build lr
-        optim_name = optim_item.keys()[0]  # get optim_name1
+        optim_name = list(optim_item.keys())[0]  # get optim_name1
         optim_scope = optim_item[optim_name].pop('scope')  # get scope
         optim_cfg = optim_item[optim_name]  # get optim_cfg
 
diff --git a/ppcls/utils/save_load.py b/ppcls/utils/save_load.py
index ca0e6bd9da32c4924a0da96dccf5ecacb93b72f8..093255379cd35875fbaf06282e391017bf7f14a3 100644
--- a/ppcls/utils/save_load.py
+++ b/ppcls/utils/save_load.py
@@ -103,7 +103,7 @@ def init_model(config, net, optimizer=None, loss: paddle.nn.Layer=None):
         metric_dict = paddle.load(checkpoints + ".pdstates")
         # set state dict
         net.set_state_dict(para_dict)
-        loss.loss_func[i].set_state_dict(para_dict)
+        loss.set_state_dict(para_dict)
         for i in range(len(optimizer)):
             optimizer[i].set_state_dict(opti_dict)
         logger.info("Finish load checkpoints from {}".format(checkpoints))