hotfix

Guitaricet · Sep 28, 2023 · 7f43aab · 7f43aab
1 parent 7bce45b
commit 7f43aab
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/peft_pretraining/args_utils.py b/peft_pretraining/args_utils.py
@@ -76,4 +76,11 @@ def check_args_torchrun_main(args):
     if args.distributed_type == "fsdp" and "zero" in args.optimizer:
         raise ValueError("FSDP does zero-optimization by default, do not specify optimizer as zero optimizer.")
 
+    if args.skip_batches is not None:
+        args.skip_batches = map(int, args.skip_batches.split(","))
+        args.skip_batches = set(args.skip_batches)
+        logger.info(f"Skipping batches {args.skip_batches}")
+
+    args.skip_batches = args.skip_batches or set()
+
     return args