diff --git a/pyext/src/restraint.py b/pyext/src/restraint.py
index d44b46e8d061c4fac40d17334cc7be1673f6b12a..6c8a2013494cebf5154aaefb4b70d2a42e5aa996 100755
--- a/pyext/src/restraint.py
+++ b/pyext/src/restraint.py
@@ -11,6 +11,7 @@ import IMP.container
 import IMP.bayesianem
 import IMP.isd
 import IMP.pmi.tools
+import IMP.pmi.mmcif
 import IMP.isd.gmm_tools
 import sys
 from math import sqrt
@@ -96,6 +97,7 @@ class GaussianEMRestraintWrapper(object):
         print('will scale target mass by', target_mass_scale)
 
         if target_fn != '':
+            self._set_dataset(target_fn)
             self.target_ps = []
             IMP.isd.gmm_tools.decorate_gmm_from_text(
                 target_fn,
@@ -117,6 +119,9 @@ class GaussianEMRestraintWrapper(object):
                 ms=IMP.atom.Mass(p).get_mass()
                 IMP.atom.Mass(p).set_mass(ms*scale)
 
+        for p, state in IMP.pmi.tools._all_protocol_outputs([], densities[0]):
+            p.add_em3d_restraint(state, self.target_ps, self.densities, self)
+
         # setup model GMM
         self.model_ps = []
         for h in self.densities:
@@ -182,6 +187,11 @@ class GaussianEMRestraintWrapper(object):
         self.rs.add_restraint(self.gaussianEM_restraint)
         self.set_weight(weight)
 
+    def _set_dataset(self, target_fn):
+        """Set the dataset to point to the input file"""
+        p = IMP.pmi.mmcif.GMMParser()
+        self.dataset = p.parse_file(target_fn)['dataset']
+
     def center_target_density_on_model(self):
         '''
         aligns the center of mass of the target GMM on the center of mass of the model