diff --git a/halo_tasks.py b/halo_tasks.py
index d01a98b3..1d7bb195 100644
--- a/halo_tasks.py
+++ b/halo_tasks.py
@@ -239,7 +239,7 @@ def process_halos(
     t0_all = time.time()
 
     # Count halos to do
-    nr_halos_left = comm.allreduce(np.sum(halo_arrays["done"].local == 0))
+    nr_halos_left = comm.allreduce(np.sum(halo_arrays["done"].local.value == 0))
 
     # Loop until all halos are done
     nr_halos = len(halo_arrays["index"].full)
@@ -262,7 +262,7 @@ def process_halos(
             t0_task = time.time()
 
             # Skip halos we already did
-            if halo_arrays["done"].full[task_to_do] == 0:
+            if halo_arrays["done"].full[task_to_do].value == 0:
 
                 # Extract this halo's VR information (centre, radius, index etc)
                 input_halo = {}
@@ -312,7 +312,7 @@ def process_halos(
 
     # Count halos left to do
     comm.barrier()
-    nr_halos_left = comm.allreduce(np.sum(halo_arrays["done"].local == 0))
+    nr_halos_left = comm.allreduce(np.sum(halo_arrays["done"].local.value == 0))
 
     # Stop the clock
     comm.barrier()
diff --git a/mask_cells.py b/mask_cells.py
index 27d33d97..fac66c57 100644
--- a/mask_cells.py
+++ b/mask_cells.py
@@ -27,7 +27,7 @@ def mask_cells(comm, cellgrid, centre, radius, done):
     for halo_nr in range(first_halo, first_halo+nr_halos_local):
         
         # Flag cells around this halo
-        if done[halo_nr] == 0:
+        if done[halo_nr].value == 0:
             pos_min = centre[halo_nr,:] - radius[halo_nr] - 0.5*cellgrid.cell_size
             pos_max = centre[halo_nr,:] + radius[halo_nr] + 0.5*cellgrid.cell_size
             cellgrid.mask_region(mask, pos_min, pos_max)