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 79abe540936b9dd54ac04a935059e784d3fea153..51665337201bcd8e1edb8169f7b38eb01287d6ba 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,16 +70,14 @@ Loss:
       mode: "l2"
       model_name_pairs:
         - ["Student", "Teacher"]
-      key: hidden_states
-      index: 5
+      key: hidden_states_5
       name: "loss_5"
   - DistillationVQADistanceLoss:
       weight: 0.5
       mode: "l2"
       model_name_pairs:
         - ["Student", "Teacher"]
-      key: hidden_states
-      index: 8
+      key: hidden_states_8
       name: "loss_8"
   
   
@@ -182,4 +180,3 @@ Eval:
     drop_last: False
     batch_size_per_card: 8
     num_workers: 4
-