diff --git a/src/engine.c b/src/engine.c
index b68b456ee24bcdd7a4b37c30bc49987fcd852691..09a00640eda2ad35c24571a43481d6ef1fe478c9 100644
--- a/src/engine.c
+++ b/src/engine.c
@@ -3397,12 +3397,34 @@ void engine_do_drift_all_mapper(void *map_data, int num_elements,
                                 void *extra_data) {
 
   const struct engine *e = (const struct engine *)extra_data;
+  const int restarting = e->restarting;
   struct space *s = e->s;
-  struct cell *cells_top = s->cells_top;
-  int *local_cells = (int *)map_data;
+  struct cell *cells_top;
+  int *local_cells_with_tasks_top;
+
+  if (restarting) {
+
+    /* When restarting, we loop over all top-level cells */
+    cells_top = (struct cell *)map_data;
+    local_cells_with_tasks_top = NULL;
+    
+  } else {
+
+    /* In any other case, we use th list of local cells with tasks */
+    cells_top = s->cells_top;
+    local_cells_with_tasks_top = (int *)map_data;
+  }
 
   for (int ind = 0; ind < num_elements; ind++) {
-    struct cell *c = &cells_top[local_cells[ind]];
+
+    struct cell *c;
+
+    /* When restarting, the list of local cells with tasks does not
+       yet exist. We use the raw list of top-level cells instead */
+    if (restarting)
+      c = &cells_top[ind];
+    else
+      c = &cells_top[local_cells_with_tasks_top[ind]];
 
     if (c->nodeID == e->nodeID) {
 
@@ -3441,9 +3463,19 @@ void engine_drift_all(struct engine *e) {
   }
 #endif
 
-  threadpool_map(&e->threadpool, engine_do_drift_all_mapper,
-                 e->s->local_cells_with_tasks_top,
-                 e->s->nr_local_cells_with_tasks, sizeof(int), 0, e);
+  if (!e->restarting) {
+
+    /* Normal case: We have a list of local cells with tasks to play with */
+    threadpool_map(&e->threadpool, engine_do_drift_all_mapper,
+                   e->s->local_cells_with_tasks_top,
+                   e->s->nr_local_cells_with_tasks, sizeof(int), 0, e);
+  } else {
+
+    /* When restarting, the list of local cells with tasks does not yet
+    /* exist. We use the raw list of top-level cells instead */
+    threadpool_map(&e->threadpool, engine_do_drift_all_mapper, e->s->cells_top,
+                   e->s->nr_cells, sizeof(struct cell), 0, e);
+  }
 
   /* Synchronize particle positions */
   space_synchronize_particle_positions(e->s);
diff --git a/src/space.c b/src/space.c
index 068fd4cc0a9218e3fdd2df48264a9ea5e2daa76a..bda2f4882588d2c183d5e013077723b30fb81512 100644
--- a/src/space.c
+++ b/src/space.c
@@ -3727,6 +3727,8 @@ void space_struct_restore(struct space *s, FILE *stream) {
   s->cells_with_particles_top = NULL;
   s->local_cells_with_particles_top = NULL;
   s->grav_top_level = NULL;
+  s->nr_local_cells_with_tasks = 0;
+  s->nr_cells_with_particles = 0;
 #ifdef WITH_MPI
   s->parts_foreign = NULL;
   s->size_parts_foreign = 0;