diff --git a/src/lerobot/datasets/factory.py b/src/lerobot/datasets/factory.py index 9f1fe6530..0dd59748d 100644 --- a/src/lerobot/datasets/factory.py +++ b/src/lerobot/datasets/factory.py @@ -127,6 +127,8 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas if cfg.dataset.use_imagenet_stats: for key in dataset.meta.camera_keys: + if key in dataset.meta.depth_keys: + continue # Exclude depth keys from ImageNet stats for stats_type, stats in IMAGENET_STATS.items(): dataset.meta.stats[key][stats_type] = torch.tensor(stats, dtype=torch.float32)