diff --git a/generate_data.py b/generate_data.py
index efca04e96..f8cc0f4f9 100644
--- a/generate_data.py
+++ b/generate_data.py
@@ -1,5 +1,6 @@
 import sys
-from estimator_new import *
+sys.path.insert(1, '../lattice-estimator')
+from estimator import *
 from sage.all import oo, save, load, ceil
 from math import log2
 import multiprocessing