diff --git a/diagonals/ohc.py b/diagonals/ohc.py
index e5cd6c63fb0951bf496a00a9c7e20a8f318cb1a3..c73298ae04030aa96ebfd6e09adbdd7063c34008 100644
--- a/diagonals/ohc.py
+++ b/diagonals/ohc.py
@@ -1,4 +1,5 @@
 import numpy as np
+import dask.array as da
 
 from numba import vectorize
 from numba import cuda
@@ -162,7 +163,7 @@ def _compute_ohc_cpu(layers, thetao, weights, area):
     """
     ohc = []
     for layer in range(len(layers)):
-        ohc_layer = _sum_red.reduce(
+        ohc_layer = da.sum(
             _multiply_array(thetao, weights[layer]),
             axis=1
         )
@@ -259,8 +260,6 @@ def _calculate_weight_numba_cuda(min_depth, max_depth, e3t, depth, mask):
     weight: float32
         Masked array containing the weights for a given layer.
     """
-    if not mask:
-        return 0
     top = depth
     bottom = top + e3t
     if bottom < min_depth or top > max_depth:
@@ -270,6 +269,8 @@ def _calculate_weight_numba_cuda(min_depth, max_depth, e3t, depth, mask):
             top = min_depth
         if bottom > max_depth:
             bottom = max_depth
+        if not mask:
+            return np.nan
 
         return (bottom - top) * 1020 * 4000
 
@@ -361,8 +362,6 @@ def _calculate_weight_numba(min_depth, max_depth, e3t, depth, mask):
     weight: float32
         Masked array containing the weights for a given layer.
     """
-    if not mask:
-        return 0
     top = depth
     bottom = top + e3t
     if bottom < min_depth or top > max_depth:
@@ -372,6 +371,8 @@ def _calculate_weight_numba(min_depth, max_depth, e3t, depth, mask):
             top = min_depth
         if bottom > max_depth:
             bottom = max_depth
+        if not mask:
+            return np.nan
 
         return (bottom - top) * 1020 * 4000
 
diff --git a/setup.py b/setup.py
index 2d084008dad0da81273c741911f8338d85652e6f..48cb8a40253d6294fe70f5283b1d6c0c005b6933 100644
--- a/setup.py
+++ b/setup.py
@@ -36,7 +36,7 @@ REQUIREMENTS = {
 
 
 setup(name='diagonals',
-      version='0.3.1',
+      version='0.3.2',
       description='Compute diagnostics targeting the CPU or the GPU',
       url='https://earth.bsc.es/gitlab/es/diagonals',
       author='BSC-CNS Earth Sciences Department',
diff --git a/test/unit/test_ohc.py b/test/unit/test_ohc.py
index a7641dfd58804a2580ef7babd3b8132130f90779..fd885cd6a13defa0fce955af070b7cdf98d27a01 100644
--- a/test/unit/test_ohc.py
+++ b/test/unit/test_ohc.py
@@ -23,7 +23,7 @@ class TestOhc(unittest.TestCase):
             cell_height=np.array(((1., 2.), (4., 3.)), dtype=np.float32),
             cell_top_depth=np.array(((0., 0.), (1., 2.)), dtype=np.float32)
         )
-        self.assertTrue(np.all(weights == np.zeros((2, 2), dtype=np.float32)))
+        self.assertTrue(np.all(np.isnan(weights)))
 
     def test_weights_mul_layers(self):
         weights = ohc.get_weights(