From 992fb177c3af3cdd9fe110b0e65f61f6f1837b20 Mon Sep 17 00:00:00 2001 From: Michel Aractingi Date: Wed, 3 Sep 2025 18:49:30 +0200 Subject: [PATCH] further memory optimizations needed due to calling `pd.concat` --- src/lerobot/datasets/lerobot_dataset.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py index 2bd277b3b..101eb20e7 100644 --- a/src/lerobot/datasets/lerobot_dataset.py +++ b/src/lerobot/datasets/lerobot_dataset.py @@ -307,6 +307,10 @@ class LeRobotDatasetMetadata: latest_df = pd.read_parquet(latest_path) df = pd.concat([latest_df, df], ignore_index=True) + # Memort optimization + del latest_df + gc.collect() + # Write the resulting dataframe from RAM to disk path = self.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx) path.parent.mkdir(parents=True, exist_ok=True) @@ -1048,6 +1052,10 @@ class LeRobotDataset(torch.utils.data.Dataset): latest_df = pd.read_parquet(latest_path) df = pd.concat([latest_df, df], ignore_index=True) + # Memort optimization + del latest_df + gc.collect() + # Write the resulting dataframe from RAM to disk path = self.root / self.meta.data_path.format(chunk_index=chunk_idx, file_index=file_idx) path.parent.mkdir(parents=True, exist_ok=True)