diff --git a/docs/source/customization.mdx b/docs/source/customization.mdx
index 7fc9211e11..baa2876836 100644
--- a/docs/source/customization.mdx
+++ b/docs/source/customization.mdx
@@ -134,7 +134,7 @@ Read more about 8-bit model loading in `transformers` [here](https://huggingface
 
 ```python
 from datasets import load_dataset
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from trl import DPOConfig, DPOTrainer
 
 model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")