fix: cast learning_rate to float (YAML parses 2e-4 as string)

This commit is contained in:
Christian Medina
2026-07-01 11:19:00 -04:00
parent 2fa52cd9d4
commit 9471d5cc77

View File

@@ -96,7 +96,7 @@ def train(config_path):
num_train_epochs=config["train_params"]["num_train_epochs"],
per_device_train_batch_size=config["train_params"]["per_device_train_batch_size"],
gradient_accumulation_steps=config["train_params"]["gradient_accumulation_steps"],
learning_rate=config["train_params"]["learning_rate"],
learning_rate=float(config["train_params"]["learning_rate"]),
lr_scheduler_type=config["train_params"]["lr_scheduler_type"],
weight_decay=config["train_params"]["weight_decay"],
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"]),