diff --git a/docker/requirements.txt b/docker/requirements.txt
index 286a6ad..7da2769 100644
--- a/docker/requirements.txt
+++ b/docker/requirements.txt
@@ -11,7 +11,7 @@ odc-dscache>=0.2.3
 odc-stac @ git+https://github.com/opendatacube/odc-stac@69bdf64
 
 # odc-stac is in PyPI
-odc-stats[ows] @ git+https://github.com/opendatacube/odc-stats@f1f2771
+odc-stats[ows] @ git+https://github.com/opendatacube/odc-stats@b404842
 
 # For ML
 tflite-runtime
diff --git a/odc/stats/plugins/lc_ml_treelite.py b/odc/stats/plugins/lc_ml_treelite.py
index dcf1c52..51c4241 100644
--- a/odc/stats/plugins/lc_ml_treelite.py
+++ b/odc/stats/plugins/lc_ml_treelite.py
@@ -139,11 +139,11 @@ def preprocess_predict_input(self, xx: xr.Dataset):
                 images += [image]
             else:
                 veg_mask = expr_eval(
-                    "where(a==_v, 1, 0)",
+                    "where((a+0.5>_v)&(a-0.5<_v), 1, 0)",
                     {"a": image},
                     name="make_mask",
                     dtype="float32",
-                    **{"_v": int(self.mask_bands[var])},
+                    **{"_v": self.mask_bands[var]},
                 )
 
         if veg_mask is None:
diff --git a/odc/stats/plugins/lc_treelite_cultivated.py b/odc/stats/plugins/lc_treelite_cultivated.py
index 7327230..37825e0 100644
--- a/odc/stats/plugins/lc_treelite_cultivated.py
+++ b/odc/stats/plugins/lc_treelite_cultivated.py
@@ -6,10 +6,10 @@
 import numpy as np
 import xarray as xr
 import dask.array as da
-import numexpr as ne
 
 from ._registry import register
 from .lc_ml_treelite import StatsMLTree, mask_and_predict
+import numexpr as ne
 
 NODATA = 255