diff --git a/strata_trapper.m b/strata_trapper.m
index 014bed2..5625d61 100644
--- a/strata_trapper.m
+++ b/strata_trapper.m
@@ -17,6 +17,7 @@
 krg = zeros(G.cells.num,3,length(saturations));
 
 cells_num = min(length(mask),G.cells.num);
+mask = mask(1:cells_num);
 
 if enable_waitbar
     wb_queue = parallel.pool.DataQueue;