diff --git a/src/scheduler.c b/src/scheduler.c
index cb8d9c4d45b89e85de6ae3efde4f0a321697fb49..641f157f0c010adf20cd50d4cd7730f215f33aa7 100644
--- a/src/scheduler.c
+++ b/src/scheduler.c
@@ -747,15 +747,25 @@ struct task *scheduler_addtask(struct scheduler *s, enum task_types type,
 void scheduler_set_unlocks(struct scheduler *s) {
 
   /* Store the counts for each task. */
-  int *counts;
-  if ((counts = (int *)malloc(sizeof(int) * s->nr_tasks)) == NULL)
+  short int *counts;
+  if ((counts = (short int *)malloc(sizeof(short int) * s->nr_tasks)) == NULL)
     error("Failed to allocate temporary counts array.");
-  bzero(counts, sizeof(int) * s->nr_tasks);
-  for (int k = 0; k < s->nr_unlocks; k++) counts[s->unlock_ind[k]] += 1;
+  bzero(counts, sizeof(short int) * s->nr_tasks);
+  for (int k = 0; k < s->nr_unlocks; k++) {
+    counts[s->unlock_ind[k]] += 1;
+
+#ifdef SWIFT_DEBUG_CHECKS
+    /* Check that we are not overflowing */
+    if (counts[s->unlock_ind[k]] < 0)
+      error("Task unlocking more than %d other tasks!",
+            (1 << (sizeof(short int) - 1)) - 1);
+#endif
+  }
 
   /* Compute the offset for each unlock block. */
-  int *offsets;
-  if ((offsets = (int *)malloc(sizeof(int) * (s->nr_tasks + 1))) == NULL)
+  short int *offsets;
+  if ((offsets = (short int *)malloc(sizeof(short int) * (s->nr_tasks + 1))) ==
+      NULL)
     error("Failed to allocate temporary offsets array.");
   offsets[0] = 0;
   for (int k = 0; k < s->nr_tasks; k++) offsets[k + 1] = offsets[k] + counts[k];
@@ -1031,10 +1041,17 @@ void scheduler_rewait_mapper(void *map_data, int num_elements,
     /* Increment the task's own wait counter for the enqueueing. */
     atomic_inc(&t->wait);
 
+#ifdef SWIFT_DEBUG_CHECKS
+    /* Check that we don't have more waits that what can be stored. */
+    if (t->wait < 0)
+      error("Task unlocked by more than %d tasks!",
+            (1 << (8 * sizeof(t->wait) - 1)) - 1);
+
     /* Skip sort tasks that have already been performed */
     if (t->type == task_type_sort && t->flags == 0) {
       error("Empty sort task encountered.");
     }
+#endif
 
     /* Sets the waits of the dependances */
     for (int k = 0; k < t->nr_unlock_tasks; k++) {