add

hpcaitech · Sep 19, 2023 · 1c36ad9 · 1c36ad9
1 parent 0ddcfb1
commit 1c36ad9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/tests/kit/model_zoo/torchrec/__init__.py b/tests/kit/model_zoo/torchrec/__init__.py
@@ -1 +1 @@
-from .torchrec import *
+#from .torchrec import *
diff --git a/tests/kit/model_zoo/transformers/chatglm2.py b/tests/kit/model_zoo/transformers/chatglm2.py
@@ -36,10 +36,11 @@ def data_gen_for_conditional_generation():
 
 config = ChatGLMConfig(num_layers=2,
                        padded_vocab_size=65024,
-                       hidden_size=1024,
-                       num_attention_heads=8,
+                       hidden_size=64,
+                       num_attention_heads=4,
                        multi_query_attention=True,
                        multi_query_group_num=2,
+                       kv_channels=16,
                        rmsnorm=True,
                        original_rope=True,
                        use_cache=True,