Update src/datachain/toolkit/split.py

iterative · Dec 10, 2024 · 24cd3b5 · 24cd3b5
1 parent bc1ea48
commit 24cd3b5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/datachain/toolkit/split.py b/src/datachain/toolkit/split.py
@@ -77,8 +77,8 @@ def train_test_split(
 
     return [
         dc.filter(
-            rand_col >= round(sum(weights_normalized[:index]) * (RESOLUTION + 1)),
-            rand_col < round(sum(weights_normalized[: index + 1]) * (RESOLUTION + 1)),
+            rand_col >= round(sum(weights_normalized[:index]) * (RESOLUTION - 1)),
+            rand_col < round(sum(weights_normalized[: index + 1]) * (RESOLUTION - 1)),
         )
         for index, _ in enumerate(weights_normalized)
     ]