[test] fix chatglm test kit (#5793)

tests/kit/model_zoo/transformers/chatglm2.py

@@ -33,41 +33,40 @@ loss_fn_for_chatglm_model = lambda x: torch.nn.functional.mse_loss(
 )
 loss_fn = lambda x: x["loss"]
 
-config = AutoConfig.from_pretrained(
+
+infer_config = AutoConfig.from_pretrained(
     "THUDM/chatglm2-6b",
     trust_remote_code=True,
     num_layers=2,
     padded_vocab_size=65024,
-    hidden_size=64,
+    hidden_size=128,
-    ffn_hidden_size=214,
     num_attention_heads=8,
+    multi_query_attention=True,
+    multi_query_group_num=2,
     kv_channels=16,
     rmsnorm=True,
     original_rope=True,
     use_cache=True,
-    multi_query_attention=False,
     torch_dtype=torch.float32,
 )
 
 
-infer_config = AutoConfig.from_pretrained(
+def init_chatglm():
+    config = AutoConfig.from_pretrained(
         "THUDM/chatglm2-6b",
         trust_remote_code=True,
         num_layers=2,
         padded_vocab_size=65024,
-    hidden_size=128,
+        hidden_size=64,
+        ffn_hidden_size=214,
         num_attention_heads=8,
-    multi_query_attention=True,
-    multi_query_group_num=2,
         kv_channels=16,
         rmsnorm=True,
         original_rope=True,
         use_cache=True,
+        multi_query_attention=False,
         torch_dtype=torch.float32,
     )
-
-
-def init_chatglm():
     model = AutoModelForCausalLM.from_config(config, empty_init=False, trust_remote_code=True)
     for m in model.modules():
         if m.__class__.__name__ == "RMSNorm":