diff --git a/gemma-peft.md b/gemma-peft.md
index 8ee4615304..a7493fd1cf 100644
--- a/gemma-peft.md
+++ b/gemma-peft.md
@@ -63,6 +63,7 @@ We start by downloading the model and the tokenizer. We also include a `BitsAndB
 
 ```python
 import torch
+import os
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
 model_id = "google/gemma-2b"