diff --git a/src/alignment/model_utils.py b/src/alignment/model_utils.py
index 52bf00f..d538fb2 100644
--- a/src/alignment/model_utils.py
+++ b/src/alignment/model_utils.py
@@ -39,11 +39,15 @@ def get_kbit_device_map() -> Dict[str, int] | None:
     return {"": get_current_device()} if torch.cuda.is_available() else None
 
 
-def get_quantization_config(model_args) -> BitsAndBytesConfig | None:
+def get_quantization_config(model_args: ModelArguments) -> BitsAndBytesConfig | None:
     if model_args.load_in_4bit:
+        compute_dtype = torch.float16
+        if model_args.torch_dtype not in {"auto", None}:
+            compute_dtype = getattr(torch, model_args.torch_dtype)
+
         quantization_config = BitsAndBytesConfig(
             load_in_4bit=True,
-            bnb_4bit_compute_dtype=torch.float16,  # For consistency with model weights, we use the same value as `torch_dtype` which is float16 for PEFT models
+            bnb_4bit_compute_dtype=compute_dtype,
             bnb_4bit_quant_type=model_args.bnb_4bit_quant_type,
             bnb_4bit_use_double_quant=model_args.use_bnb_nested_quant,
         )