diff --git a/meson.build b/meson.build
index e6004f4ee..126edad4a 100644
--- a/meson.build
+++ b/meson.build
@@ -4,7 +4,7 @@ project(
   # Note that the git commit hash cannot be added dynamically here
   # That only happens when importing from a git repository.
   # See `sktree/__init__.py`
-  version: '0.3.0dev0',
+  version: '0.2.1',
   license: 'BSD-3',
   meson_version: '>= 0.64.0',
   default_options: [
diff --git a/pyproject.toml b/pyproject.toml
index d9d901e32..c814a7966 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ name = "scikit-tree"
 # 0.0.0 is standard placeholder for poetry-dynamic-versioning
 # any changes to this should not be checked in
 #
-version = "0.3.0dev0"
+version = "0.2.1"
 description = "Modern decision trees in Python"
 maintainers = [
     {name="Neurodata", email="adam.li@columbia.edu"}
diff --git a/sktree/__init__.py b/sktree/__init__.py
index 1aa58557f..05d167a3f 100644
--- a/sktree/__init__.py
+++ b/sktree/__init__.py
@@ -3,7 +3,7 @@
 import os
 import sys
 
-__version__ = "0.3.0dev0"
+__version__ = "0.2.1"
 logger = logging.getLogger(__name__)
 
 
diff --git a/sktree/stats/tests/test_forestht.py b/sktree/stats/tests/test_forestht.py
index bf1c75d0d..f63894be7 100644
--- a/sktree/stats/tests/test_forestht.py
+++ b/sktree/stats/tests/test_forestht.py
@@ -314,7 +314,7 @@ def test_sample_size_consistency_of_estimator_indices_(permute_per_tree, sample_
     )
     if sample_dataset_per_tree:
         assert_array_equal(
-            samples,
+            sorted(np.unique(samples)),
             sorted(np.unique(np.concatenate([x[1] for x in clf.train_test_samples_]).flatten())),
         )
     else: