diff --git a/src/timestep.h b/src/timestep.h
index 63a21bf3e2c82344f9a0b1ac5f4ac769af10c7d1..b3ef96493772e7e93853ad28a0536e4f8448dceb 100644
--- a/src/timestep.h
+++ b/src/timestep.h
@@ -70,14 +70,15 @@ make_integer_timestep(float new_dt, timebin_t old_bin, integertime_t ti_current,
 __attribute__((always_inline)) INLINE static integertime_t get_gpart_timestep(
     const struct gpart *restrict gp, const struct engine *restrict e) {
 
-  const float new_dt_external = external_gravity_timestep(
-      e->time, e->external_potential, e->physical_constants, gp);
+  float new_dt = FLT_MAX;
 
-  /* const float new_dt_self = */
-  /*     gravity_compute_timestep_self(e->physical_constants, gp); */
-  const float new_dt_self = FLT_MAX;  // MATTHIEU
+  if (e->policy & engine_policy_external_gravity)
+    new_dt =
+        min(new_dt, external_gravity_timestep(e->time, e->external_potential,
+                                              e->physical_constants, gp));
 
-  float new_dt = min(new_dt_external, new_dt_self);
+  if (e->policy & engine_policy_self_gravity)
+    new_dt = min(new_dt, gravity_compute_timestep_self(gp));
 
   /* Limit timestep within the allowed range */
   new_dt = min(new_dt, e->dt_max);
@@ -114,14 +115,13 @@ __attribute__((always_inline)) INLINE static integertime_t get_part_timestep(
   float new_dt_grav = FLT_MAX;
   if (p->gpart != NULL) {
 
-    const float new_dt_external = external_gravity_timestep(
-        e->time, e->external_potential, e->physical_constants, p->gpart);
+    if (e->policy & engine_policy_external_gravity)
+      new_dt_grav = min(new_dt_grav, external_gravity_timestep(
+                                         e->time, e->external_potential,
+                                         e->physical_constants, p->gpart));
 
-    /* const float new_dt_self = */
-    /*     gravity_compute_timestep_self(e->physical_constants, p->gpart); */
-    const float new_dt_self = FLT_MAX;  // MATTHIEU
-
-    new_dt_grav = min(new_dt_external, new_dt_self);
+    if (e->policy & engine_policy_self_gravity)
+      new_dt_grav = min(new_dt_grav, gravity_compute_timestep_self(p->gpart));
   }
 
   /* Final time-step is minimum of hydro and gravity */