diff --git a/aiida_nanotech_empa/workflows/cp2k/reftraj_md_workchain.py b/aiida_nanotech_empa/workflows/cp2k/reftraj_md_workchain.py
index 8a2f7e6..c6fb787 100644
--- a/aiida_nanotech_empa/workflows/cp2k/reftraj_md_workchain.py
+++ b/aiida_nanotech_empa/workflows/cp2k/reftraj_md_workchain.py
@@ -169,7 +169,7 @@ def run_reftraj_batches(self):
 
     def merge_batches_output(self):
         """Merge the output of the succefull batches only."""
-        
+
         # merged_traj = []
         # for i_batch in range(self.ctx.n_batches):
         #    merged_traj.extend(self.ctx[f"reftraj_batch_{i_batch}"].outputs.trajectory)
@@ -178,8 +178,7 @@ def merge_batches_output(self):
         forces=[self.ctx.first_structure.outputs.output_trajectory.get_array('forces')]
         for batch in self.ctx.batches:
             key = f"reftraj_batch_{batch[0]}_to_{batch[-1]}"
-            if not  getattr(
-            self.ctx, key).is_finished_ok:
+            if not getattr(self.ctx, key).is_finished_ok:
                 self.report(f"Batch {key} failed")
                 return self.exit_codes.ERROR_TERMINATION
             positions.append(getattr(self.ctx, key).outputs.output_trajectory.get_array('positions'))
@@ -192,7 +191,6 @@ def merge_batches_output(self):
         symbols = self.ctx.first_structure.outputs.output_trajectory.symbols
         output_trajectory = TrajectoryData()   
         output_trajectory.set_trajectory(symbols, positions, cells=cells) 
-        self.out('output_trajectory', self.inputs.trajectory)
+        self.out('output_trajectory', output_trajectory)
         self.report("done")
         return engine.ExitCode(0)
-