diff --git a/src/eynollah/training/inference.py b/src/eynollah/training/inference.py
index 2b26210..c38b79f 100644
--- a/src/eynollah/training/inference.py
+++ b/src/eynollah/training/inference.py
@@ -25,7 +25,7 @@ from .gt_gen_utils import (
     resize_image,
     update_list_and_return_first_with_length_bigger_than_one
 )
-from .models import (
+from ..patch_encoder import (
     PatchEncoder,
     Patches
 )
diff --git a/src/eynollah/training/models.py b/src/eynollah/training/models.py
index 115a196..6182c9e 100644
--- a/src/eynollah/training/models.py
+++ b/src/eynollah/training/models.py
@@ -30,7 +30,7 @@ from tensorflow.keras.layers import (
 from tensorflow.keras.models import Model
 from tensorflow.keras.regularizers import l2
 
-from eynollah.patch_encoder import Patches, PatchEncoder
+from ..patch_encoder import Patches, PatchEncoder
 
 ##mlp_head_units = [512, 256]#[2048, 1024]
 ###projection_dim = 64
diff --git a/src/eynollah/training/weights_ensembling.py b/src/eynollah/training/weights_ensembling.py
index 01532fd..e3ede24 100644
--- a/src/eynollah/training/weights_ensembling.py
+++ b/src/eynollah/training/weights_ensembling.py
@@ -12,7 +12,7 @@ tf_disable_interactive_logs()
 import tensorflow as tf
 from tensorflow.keras.models import load_model
 
-from .models import (
+from ..patch_encoder import (
     PatchEncoder,
     Patches,
 )