Add hidden hidden size args

f9ce8e09 · Sengxian · 4eec9807 · f9ce8e09
Commit f9ce8e09 authored May 31, 2021 by Sengxian
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 0 deletions

fmoe/megatron/utils.py fmoe/megatron/utils.py +1 -0

No files found.
--- a/fmoe/megatron/utils.py
+++ b/fmoe/megatron/utils.py
@@ -9,5 +9,6 @@ def add_fmoe_args(parser):
    group.add_argument("--top-k", type=int, default=2)
    group.add_argument("--balance-loss-weight", type=float, default=1)
    group.add_argument("--balance-strategy", type=str, default=None)
+    group.add_argument("--hidden-hidden-size", type=int, default=None)

    return parser