From 6a8be97bb59562f9fd79adcc468eb1ed1c002b57 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Tue, 11 Mar 2025 13:29:55 +0100 Subject: [PATCH] remove power of 2 optimization --- lerobot/common/datasets/sampler.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/lerobot/common/datasets/sampler.py b/lerobot/common/datasets/sampler.py index 941a5561..08c3ac13 100644 --- a/lerobot/common/datasets/sampler.py +++ b/lerobot/common/datasets/sampler.py @@ -72,14 +72,11 @@ class SumTree: def __init__(self, capacity: int): """ Args: - capacity: Maximum number of elements. The tree size is the next power of 2 for efficiency. + capacity: Maximum number of elements. """ self.capacity = capacity - self.size = 1 - while self.size < capacity: - self.size *= 2 # Ensure power-of-two size for efficient updates - - self.tree = [0.0] * (2 * self.size) # Tree structure + self.size = capacity + self.tree = [0.0] * (2 * self.size) def initialize_tree(self, priorities: List[float]): """