diff --git a/configs/kie/vi_layoutxlm/ser_vi_layoutxlm_xfund_zh_udml.yml b/configs/kie/vi_layoutxlm/ser_vi_layoutxlm_xfund_zh_udml.yml
index 238bbd2b2c7083b5534062afd3e6c11a87494a56..79abe540936b9dd54ac04a935059e784d3fea153 100644
--- a/configs/kie/vi_layoutxlm/ser_vi_layoutxlm_xfund_zh_udml.yml
+++ b/configs/kie/vi_layoutxlm/ser_vi_layoutxlm_xfund_zh_udml.yml
@@ -70,14 +70,16 @@ Loss:
       mode: "l2"
       model_name_pairs:
         - ["Student", "Teacher"]
-      key: hidden_states_5
+      key: hidden_states
+      index: 5
       name: "loss_5"
   - DistillationVQADistanceLoss:
       weight: 0.5
       mode: "l2"
       model_name_pairs:
         - ["Student", "Teacher"]
-      key: hidden_states_8
+      key: hidden_states
+      index: 8
       name: "loss_8"