diff --git a/trainer_pti.py b/trainer_pti.py
index 11e3b84..b8f599b 100755
--- a/trainer_pti.py
+++ b/trainer_pti.py
@@ -13,8 +13,17 @@ from peft import LoraConfig, get_peft_model
 from diffusers.optimization import get_scheduler
 from tqdm import tqdm
 
-from trainer.dataset_and_utils import *
-from trainer.utils.lora import *
+from trainer.dataset_and_utils import (
+    PreprocessedDataset, 
+    plot_torch_hist, 
+    plot_loss, 
+    plot_lrs,
+    load_models
+)
+from trainer.utils.lora import (
+    save_lora,
+    TokenEmbeddingsHandler
+)
 
 from trainer.utils.dtype import dtype_map
 from trainer.config import TrainingConfig