diff --git a/src/hashpipe.c b/src/hashpipe.c
index 9dfeabf..0c25f7b 100644
--- a/src/hashpipe.c
+++ b/src/hashpipe.c
@@ -11,7 +11,6 @@
 #include <pthread.h>
 #include <sys/types.h>
 #include <sys/socket.h>
-#include <sched.h>
 #include <signal.h>
 #include <poll.h>
 #include <getopt.h>
@@ -63,15 +62,14 @@ static void set_exit_status(hashpipe_thread_args_t *args) {
 
 // Function to set cpu affinity
 int
-set_cpu_affinity(unsigned int mask)
+set_cpu_affinity(__cpu_mask mask)
 {
     int i, rv;
 
     if(mask != 0) {
         cpu_set_t cpuset;
         CPU_ZERO(&cpuset);
-        // Only handle 32 cores (for now)
-        for(i=0; i<32; i++) {
+        for(i=0; i<__NCPUBITS; i++) {
             if(mask&1) {
               CPU_SET(i, &cpuset);
             }
@@ -456,7 +454,7 @@ int main(int argc, char *argv[])
 
         case 'c': // CPU number
           i = strtol(optarg, NULL, 0);
-          args[num_threads].cpu_mask = (1<<i);
+          args[num_threads].cpu_mask = (((__cpu_mask) 1)<<i);
           break;
 
         case 'p': // Load plugin
diff --git a/src/hashpipe.h b/src/hashpipe.h
index f364ad8..312ae87 100644
--- a/src/hashpipe.h
+++ b/src/hashpipe.h
@@ -2,6 +2,7 @@
 #define _HASHPIPE_H
 
 #include <stdio.h>
+#include <sched.h>
 
 #include "hashpipe_error.h"
 #include "hashpipe_databuf.h"
@@ -139,7 +140,7 @@ struct hashpipe_thread_args {
     int instance_id;
     int input_buffer;
     int output_buffer;
-    unsigned int cpu_mask; // 0 means use inherited
+    __cpu_mask cpu_mask; // 0 means use inherited
     int finished;
     pthread_cond_t finished_c;
     pthread_mutex_t finished_m;