diff --git a/src/common_io.c b/src/common_io.c
index e5da5838ca4654876b1ca76a2fc0e36e3611f2a8..edb116e604f0cda44c4f028a409ea6576dab32da 100644
--- a/src/common_io.c
+++ b/src/common_io.c
@@ -1357,10 +1357,12 @@ void io_collect_sparts_to_write(const struct spart* restrict sparts,
  * @param Ngparts The total number of #part.
  * @param Ngparts_written The total number of #part to write.
  */
-void io_collect_gparts_to_write(const struct gpart* restrict gparts,
-                                struct gpart* restrict gparts_written,
-                                const size_t Ngparts,
-                                const size_t Ngparts_written) {
+void io_collect_gparts_to_write(
+    const struct gpart* restrict gparts,
+    const struct velociraptor_gpart_data* restrict vr_data,
+    struct gpart* restrict gparts_written,
+    struct velociraptor_gpart_data* restrict vr_data_written,
+    const size_t Ngparts, const size_t Ngparts_written, const int with_stf) {
 
   size_t count = 0;
 
@@ -1372,6 +1374,8 @@ void io_collect_gparts_to_write(const struct gpart* restrict gparts,
         (gparts[i].time_bin != time_bin_not_created) &&
         (gparts[i].type == swift_type_dark_matter)) {
 
+      if (with_stf) vr_data_written[count] = vr_data[i];
+
       gparts_written[count] = gparts[i];
       count++;
     }
@@ -1379,7 +1383,7 @@ void io_collect_gparts_to_write(const struct gpart* restrict gparts,
 
   /* Check that everything is fine */
   if (count != Ngparts_written)
-    error("Collected the wrong number of s-particles (%zu vs. %zu expected)",
+    error("Collected the wrong number of g-particles (%zu vs. %zu expected)",
           count, Ngparts_written);
 }
 
diff --git a/src/common_io.h b/src/common_io.h
index bdff3e37d11c59a71c08dcf3e19758019adc3a73..eb1ee0a804f324d897842fb2a0ca33fc07e769d6 100644
--- a/src/common_io.h
+++ b/src/common_io.h
@@ -36,6 +36,7 @@
 struct cell;
 struct part;
 struct gpart;
+struct velociraptor_gpart_data;
 struct spart;
 struct xpart;
 struct io_props;
@@ -113,9 +114,11 @@ void io_collect_sparts_to_write(const struct spart* restrict sparts,
                                 const size_t Nsparts,
                                 const size_t Nsparts_written);
 void io_collect_gparts_to_write(const struct gpart* restrict gparts,
+                                const struct velociraptor_gpart_data* vr_data,
                                 struct gpart* restrict gparts_written,
+                                struct velociraptor_gpart_data* vr_data_written,
                                 const size_t Ngparts,
-                                const size_t Ngparts_written);
+                                const size_t Ngparts_written, int with_stf);
 void io_prepare_dm_gparts(struct threadpool* tp, struct gpart* const gparts,
                           size_t Ndm);
 void io_duplicate_hydro_gparts(struct threadpool* tp, struct part* const parts,