From 6b137f0b80bd30e3c722af81104f442b43a08c51 Mon Sep 17 00:00:00 2001 From: Cadene Date: Wed, 17 Apr 2024 20:48:13 +0000 Subject: [PATCH] small fix --- lerobot/scripts/train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lerobot/scripts/train.py b/lerobot/scripts/train.py index 83636260..6cd02f18 100644 --- a/lerobot/scripts/train.py +++ b/lerobot/scripts/train.py @@ -5,7 +5,7 @@ from pathlib import Path import hydra import torch from datasets import Dataset, concatenate_datasets -from datasets.utils.logging import disable_progress_bar, enable_progress_bars +from datasets.utils import disable_progress_bars, enable_progress_bars from lerobot.common.datasets.factory import make_dataset from lerobot.common.datasets.utils import cycle @@ -177,7 +177,7 @@ def add_episodes_inplace( example["episode_data_index_to"] += start_index return example - disable_progress_bar() # map has a tqdm progress bar + disable_progress_bars() # map has a tqdm progress bar hf_dataset = hf_dataset.map(shift_indices) enable_progress_bars()