diff --git a/lora_utils.py b/lora_utils.py
index 9ed8c6b..870dc6a 100644
--- a/lora_utils.py
+++ b/lora_utils.py
@@ -2,8 +2,9 @@ import os, json
 import torch
 from safetensors.torch import load_file
 from typing import Dict
-
+from peft import PeftModel
 from dataset_and_utils import TokenEmbeddingsHandler
+from safetensors.torch import save_file
 
 '''
 from diffusers.utils import (