fix

preprocess_all.py

@@ -72,6 +72,11 @@ if __name__ == "__main__":
         action="store_true",
         help="Use JP-Extra model",
     )
+    parser.add_argument(
+        "--matcha_only",
+        action="store_true",
+        help="Train only the Matcha Flow and Differential Attention modules",
+    )
     parser.add_argument(
         "--val_per_lang",
         type=int,
@@ -110,4 +115,5 @@ if __name__ == "__main__":
         val_per_lang=args.val_per_lang,
         log_interval=args.log_interval,
         yomi_error=args.yomi_error,
+        matcha_only=args.matcha_only,
     )