Update pretrain_bert_demo.py

Tencent · Dec 8, 2021 · 0539721 · 0539721
1 parent f5fee95
commit 0539721
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/pretrain_bert_demo.py b/examples/pretrain_bert_demo.py
@@ -592,12 +592,12 @@ def visit_and_register_hooks(module):
         SEQ_LEN = 1024
         NUM_LAYER = 65
         NUM_HEAD = 16
-    elif MODEL_NAME == "GPT3_18B":
+    elif MODEL_NAME == "GPT3_15B":
         HIDDEN_DIM = 4096
         SEQ_LEN = 1024
         NUM_LAYER = 78
         NUM_HEAD = 16
-    elif MODEL_NAME == "GPT3_17B":
+    elif MODEL_NAME == "GPT3_18B":
         HIDDEN_DIM = 4096
         SEQ_LEN = 1024
         NUM_LAYER = 90