[test] smaller gpt2 test case

2024-05-27 03:42:51 +00:00 · 2024-05-27 03:42:51 +00:00 · df705a5210
parent b96c6390f4
commit df705a5210
1 changed files with 1 additions and 1 deletions
--- a/tests/kit/model_zoo/transformers/gpt.py
+++ b/tests/kit/model_zoo/transformers/gpt.py
@ -121,7 +121,7 @@ loss_fn = lambda x: x["loss"]
 config = transformers.GPT2Config(
    n_layer=2,
    n_head=4,
-    n_embd=128,
+    n_embd=16,
    vocab_size=50258,
    attn_pdrop=0,
    embd_pdrop=0,