diff --git a/mofa/generator.py b/mofa/generator.py
index 75add995..b904050a 100644
--- a/mofa/generator.py
+++ b/mofa/generator.py
@@ -3,7 +3,7 @@
 
 import ase
 
-from model import MOFRecord
+from mofa.model import MOFRecord
 from difflinker_sample import sample_from_sdf
 from difflinker_train import get_args, main
 import yaml