diff --git a/modules/training.py b/modules/training.py
index d0018a0f..e629d86e 100644
--- a/modules/training.py
+++ b/modules/training.py
@@ -14,7 +14,7 @@ from datasets import Dataset, load_dataset
 from peft import (
     LoraConfig,
     get_peft_model,
-    prepare_model_for_kbit_training,
+    prepare_model_for_int8_training,
     set_peft_model_state_dict
 )