fix: torch_dtype -> dtype (fix deprecation warning)

2026-07-01 07:38:36 -04:00
parent f651ca030c
commit f0ee6bc9a2
1 changed files with 1 additions and 1 deletions
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -36,7 +36,7 @@ def train(config_path):
    print(f"Loading model: {config['base_model']}")
    model = AutoModelForCausalLM.from_pretrained(
        config["base_model"],
-        torch_dtype=torch.bfloat16,  # Convert FP8 -> bf16
+        dtype=torch.bfloat16,  # Convert FP8 -> bf16
        device_map="cpu",  # Load to CPU first
        trust_remote_code=True,
    )