fixing export checkpoint (other thigns remain)

RWKV · Aug 19, 2023 · 44b53ab · 44b53ab
1 parent 1bb5806
commit 44b53ab
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/RWKV-v4neo/export_checkpoint.py b/RWKV-v4neo/export_checkpoint.py
@@ -631,13 +631,13 @@ def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
         "output_file",
         type=str,
         default="",
-        required=False,
+        nargs="?",
         help="path to the pytorch state_dict output file (e.g. path/model.pth), if blank defaults to 'rwkv_model.pth' within the checkpoint dir")
     parser.add_argument(
         "dtype",
         type=str,
         default="bf16",
-        required=False,
+        nargs="?",
         help="dtype to save the model as, either 'bf16' or 'fp32'")
 
     parser.add_argument("-d", "--debug", action='store_true', help="enable debug")

diff --git a/RWKV-v5/export_checkpoint.py b/RWKV-v5/export_checkpoint.py
@@ -631,13 +631,13 @@ def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
         "output_file",
         type=str,
         default="",
-        required=False,
+        nargs="?",
         help="path to the pytorch state_dict output file (e.g. path/model.pth), if blank defaults to 'rwkv_model.pth' within the checkpoint dir")
     parser.add_argument(
         "dtype",
         type=str,
         default="bf16",
-        required=False,
+        nargs="?",
         help="dtype to save the model as, either 'bf16' or 'fp32'")
 
     parser.add_argument("-d", "--debug", action='store_true', help="enable debug")