diff --git a/src/run_modules/model_setup.py b/src/run_modules/model_setup.py
index 812f639c8f18725e1e7c4decd03d5f2b61578fc4..a47ef67ad5781ff37ce812aa931dbd195d4513dc 100644
--- a/src/run_modules/model_setup.py
+++ b/src/run_modules/model_setup.py
@@ -15,6 +15,7 @@ from src.run_modules.run_environment import RunEnvironment
 from src.helpers import l_p_loss, LearningRateDecay
 from src.model_modules.inception_model import InceptionModelBase
 from src.model_modules.flatten import flatten_tail
+from src.model_modules.model_class import MyLittleModel
 
 
 class ModelSetup(RunEnvironment):