feat: stabilize meanflow training and time sampling

2026-03-11 22:54:48 +08:00
parent 9b2968997c
commit 5897a0afd1
2 changed files with 17 additions and 11 deletions
--- a/main.py
+++ b/main.py
@@ -11,6 +11,7 @@ def build_parser() -> argparse.ArgumentParser:
    parser.add_argument("--seq-len", type=int, default=None)
    parser.add_argument("--lr", type=float, default=None)
    parser.add_argument("--weight-decay", type=float, default=None)
+    parser.add_argument("--max-grad-norm", type=float, default=None)
    parser.add_argument("--device", type=str, default=None)
    parser.add_argument("--output-dir", type=str, default=None)
    parser.add_argument("--project", type=str, default=None)