diff --git a/dpsynth/bin/run_data_generation.py b/dpsynth/bin/run_data_generation.py
index 2cca438..d9bccf4 100644
--- a/dpsynth/bin/run_data_generation.py
+++ b/dpsynth/bin/run_data_generation.py
@@ -29,6 +29,7 @@
 from dpsynth.pipeline_transformations import aim
 from dpsynth.pipeline_transformations import input_output
 from dpsynth.pipeline_transformations import types
+from google.protobuf import message
 import pipeline_dp