diff --git a/configs/keypoint/higherhrnet/higherhrnet_hrnet_w32_512.yml b/configs/keypoint/higherhrnet/higherhrnet_hrnet_w32_512.yml
index 25a0ab43cbfedacd5bcd938683280af57f86b23e..f4fcdfbea22a9514d3fa2bd8fa3a30326af6b2d7 100644
--- a/configs/keypoint/higherhrnet/higherhrnet_hrnet_w32_512.yml
+++ b/configs/keypoint/higherhrnet/higherhrnet_hrnet_w32_512.yml
@@ -1,8 +1,8 @@
 use_gpu: true
-log_iter: 1
+log_iter: 10
 save_dir: output
 snapshot_epoch: 10
-weights: output/higherhrnet_hrnet_v1_512/290
+weights: output/higherhrnet_hrnet_w32_512/model_final
 epoch: 300
 num_joints: &num_joints 17
 flip_perm: &flip_perm [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]
diff --git a/ppdet/modeling/heads/keypoint_hrhrnet_head.py b/ppdet/modeling/heads/keypoint_hrhrnet_head.py
index 90a1298fcec398cd309f8371641b7f6b0ede2a3d..2cec010bb7844dd35c309eb1e83c33a3b2c917db 100644
--- a/ppdet/modeling/heads/keypoint_hrhrnet_head.py
+++ b/ppdet/modeling/heads/keypoint_hrhrnet_head.py
@@ -92,15 +92,14 @@ class HrHRNetHead(nn.Layer):
         xo2 = self.conv2(x2)
         num_joints = self.num_joints
         if self.training:
+            heatmap1, tagmap = paddle.split(xo1, 2, axis=1)
             if self.swahr:
                 so1 = self.scalelayer0(x1)
                 so2 = self.scalelayer1(x2)
-                hrhrnet_outputs = ([xo1[:, :num_joints], so1], [xo2, so2],
-                                   xo1[:, num_joints:])
+                hrhrnet_outputs = ([heatmap1, so1], [xo2, so2], tagmap)
                 return self.loss(hrhrnet_outputs, targets)
             else:
-                hrhrnet_outputs = (xo1[:, :num_joints], xo2,
-                                   xo1[:, num_joints:])
+                hrhrnet_outputs = (heatmap1, xo2, tagmap)
                 return self.loss(hrhrnet_outputs, targets)
 
         # averaged heatmap, upsampled tagmap
diff --git a/ppdet/modeling/losses/keypoint_loss.py b/ppdet/modeling/losses/keypoint_loss.py
index 442204cc3703c4e551eeb1bcc901286b64fd2c0b..d37ef5759a9c9c18692ac8c7baae1e47051dbf62 100644
--- a/ppdet/modeling/losses/keypoint_loss.py
+++ b/ppdet/modeling/losses/keypoint_loss.py
@@ -194,7 +194,10 @@ class AELoss(object):
 
     def __call__(self, preds, tagmaps):
         bs = preds.shape[0]
-        losses = [self.apply_single(preds[i], tagmaps[i]) for i in range(bs)]
+        losses = [
+            self.apply_single(preds[i:i + 1].squeeze(),
+                              tagmaps[i:i + 1].squeeze()) for i in range(bs)
+        ]
         pull = self.pull_factor * sum(loss[0] for loss in losses) / len(losses)
         push = self.push_factor * sum(loss[1] for loss in losses) / len(losses)
         return pull, push