diff --git a/training/scripts/train.py b/training/scripts/train.py
index 43bdc6a..0bcbe3b 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -36,17 +36,18 @@ def train(config_path):
     print(f"Loading model: {config['base_model']}")
     from transformers import BitsAndBytesConfig
     
-    quantization_config = BitsAndBytesConfig(
+    bnb_config = BitsAndBytesConfig(
         load_in_4bit=True,
-        bnb_4bit_compute_dtype=torch.bfloat16,
         bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16,
         bnb_4bit_use_double_quant=True,
     )
     
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
-        quantization_config=quantization_config,
-        device_map="auto",  # Distribute across GPUs
+        quantization_config=bnb_config,
+        dtype=torch.bfloat16,
+        device_map="auto",
         trust_remote_code=True,
     )
     print("Model loaded with QLoRA (4-bit).")