diff --git a/src/engine.c b/src/engine.c
index 51392049aad1a95ebd64a36477bc7ff68ecef211..699c59c6a8110be5b557b97dfef27e30ecd4f198 100644
--- a/src/engine.c
+++ b/src/engine.c
@@ -372,7 +372,8 @@ void engine_repartition(struct engine *e) {
       /* Skip un-interesting tasks. */
       if (t->type != task_type_self && t->type != task_type_pair &&
           t->type != task_type_sub && t->type != task_type_ghost &&
-          t->type != task_type_kick1 && t->type != task_type_kick2)
+          t->type != task_type_drift && t->type != task_type_kick &&
+          t->type != task_type_init)
         continue;
 
       /* Get the task weight. */
@@ -403,8 +404,8 @@ void engine_repartition(struct engine *e) {
       int cid = ci - cells;
 
       /* Different weights for different tasks. */
-      if (t->type == task_type_ghost || t->type == task_type_kick1 ||
-          t->type == task_type_kick2) {
+      if (t->type == task_type_ghost || t->type == task_type_drift ||
+          t->type == task_type_kick) {
         /* Particle updates add only to vertex weight. */
         weights_v[cid] += w;
 
@@ -1971,7 +1972,7 @@ void engine_makeproxies(struct engine *e) {
  * @param grid The grid.
  */
 
-void engine_split(struct engine *e, int *grid) {
+void engine_split(struct engine *e, struct initpart *ipart) {
 
 #ifdef WITH_MPI
 
@@ -2115,6 +2116,7 @@ void engine_split(struct engine *e, int *grid) {
   free(s->xparts);
   s->parts = parts_new;
   s->xparts = xparts_new;
+#endif
 }
 
 #if defined(HAVE_LIBNUMA) && defined(_GNU_SOURCE)