From 1c36ad96e938d91730565b80d9b5f616416c7c3d Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 19 Sep 2023 13:54:54 +0800 Subject: [PATCH] add --- tests/kit/model_zoo/torchrec/__init__.py | 2 +- tests/kit/model_zoo/transformers/chatglm2.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/kit/model_zoo/torchrec/__init__.py b/tests/kit/model_zoo/torchrec/__init__.py index 43952e6998cf..4a19f2449602 100644 --- 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 index d74a1b29b788..c62d121db223 100644 --- 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,