diff --git a/training/scripts/train.py b/training/scripts/train.py
index fed7e20..627d90c 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -86,7 +86,7 @@ def train(config_path):
         learning_rate=config["train_params"]["learning_rate"],
         lr_scheduler_type=config["train_params"]["lr_scheduler_type"],
         weight_decay=config["train_params"]["weight_decay"],
-        warmup_ratio=config["train_params"]["warmup_ratio"],
+        warmup_steps=int(config["train_params"]["warmup_ratio"] * config["train_params"]["num_train_epochs"] * len(dataset["train"]) // config["train_params"]["per_device_train_batch_size"]),
         logging_steps=config["train_params"]["logging_steps"],
         save_steps=config["train_params"]["save_steps"],
         save_total_limit=config["train_params"]["save_total_limit"],