WIP

1 files changed, 2 insertions, 2 deletions

diff --git a/finetuning.py b/finetuning.py
index 3d7b011..b1cb7bd 100755
--- a/finetuning.py
+++ b/finetuning.py
@@ -75,10 +75,10 @@ def main():
     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)
-    parser.add_argument("--4bit", action="store_true")
+    parser.add_argument("--bnb4bit", action="store_true")
     args = parser.parse_args()
     
-    if args.4bit:
+    if args.bnb4bit:
         bnb_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_use_double_quant=True,