diff --git a/src/maggotuba/models/predict_model.py b/src/maggotuba/models/predict_model.py
index afcae80017c4cfcdb3ab11ade54d4ef34b114920..9ed347fa2ef2fdd50920d52c708b4f4826f16da9 100644
--- a/src/maggotuba/models/predict_model.py
+++ b/src/maggotuba/models/predict_model.py
@@ -3,7 +3,7 @@ from taggingbackends.data.chore import load_spine
 import taggingbackends.data.fimtrack as fimtrack
 from taggingbackends.data.labels import Labels
 from taggingbackends.features.skeleton import get_5point_spines
-from denselayer import DenseLayer
+from maggotuba.models.denselayer import DenseLayer
 import numpy as np
 import json
 
diff --git a/src/maggotuba/models/train_model.py b/src/maggotuba/models/train_model.py
index 5527b97429f2933f1afcdd178838bed557a12199..34c87f63a34307ce63a9276b66317e399f577667 100644
--- a/src/maggotuba/models/train_model.py
+++ b/src/maggotuba/models/train_model.py
@@ -1,6 +1,6 @@
 from taggingbackends.data.labels import Labels
 from taggingbackends.data.dataset import LarvaDataset
-from denselayer import DenseLayer, device
+from maggotuba.models.denselayer import DenseLayer, device
 import numpy as np
 import json
 import torch