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)