diff --git a/api/src/main.py b/api/src/main.py
index 7304312..67f00b9 100644
--- a/api/src/main.py
+++ b/api/src/main.py
@@ -1,7 +1,6 @@
 import sys
 
 sys.path.insert(1, './src/yolo')
-from yolo_model_production import convert_prediction, Yolo_model
 from train_yolo_faces import train_yolo
 
 sys.path.insert(2, './src/bird_cnn')
diff --git a/api/src/yolo/train_yolo_faces.py b/api/src/yolo/train_yolo_faces.py
index e3b20a9..63beb59 100644
--- a/api/src/yolo/train_yolo_faces.py
+++ b/api/src/yolo/train_yolo_faces.py
@@ -12,7 +12,7 @@ import torch
 from torchvision.utils import draw_bounding_boxes
 import torch.nn.functional as F
 
-from yolo_model_production import convert_prediction
+from production.yolo_model_production import convert_prediction
 
 def view_data(dataset):
     dataloader = DataLoader(dataset=dataset, batch_size=1, shuffle=False)