WIP

1 files changed, 10 insertions, 10 deletions

diff --git a/finetuning.py b/finetuning.py
index 1a44345..a46b544 100755
--- a/finetuning.py
+++ b/finetuning.py
@@ -65,16 +65,16 @@ def main():
     parser.add_argument("--model", default="google/gemma-3-4b-it")
     parser.add_argument("--optimizer", default="adamw_torch_fused")
     parser.add_argument("--gradient-checkpointing", action="store_true")
-    parser.add_argument("--batch-size", default=1)
-    parser.add_argument("--gradient-accumulation", default=4)
-    parser.add_argument("--learning-rate", default=1e-4)
-    parser.add_argument("--epochs", default=1)
-    parser.add_argument("--warmup-ratio", default=0.03)
-    parser.add_argument("--scheduler-type", default="constant")
-    parser.add_argument("--logging-steps", default=10)
-    parser.add_argument("--lora-alpha", default=32)
-    parser.add_argument("--lora-dropout", default=0.05)
-    parser.add_argument("--lora-r", default=16)
+    parser.add_argument("--batch-size", default=1, type=int)
+    parser.add_argument("--gradient-accumulation", default=4, type=int)
+    parser.add_argument("--learning-rate", default=1e-4, type=float)
+    parser.add_argument("--epochs", default=1, type=int)
+    parser.add_argument("--warmup-ratio", default=0.03, type=float)
+    parser.add_argument("--scheduler-type", default="constant", type=str)
+    parser.add_argument("--logging-steps", default=3, type=int)
+    parser.add_argument("--lora-alpha", default=32, type=int)
+    parser.add_argument("--lora-dropout", default=0.05, type=float)
+    parser.add_argument("--lora-r", default=16, type=int)
     args = parser.parse_args()
 
     bnb_config = BitsAndBytesConfig(