fix(ImageNet): excluding dpeth frames from ImageNet stats

This commit is contained in:
CarolinePascal
2026-06-24 17:34:08 +02:00
parent 5c187a06f5
commit a98a38160c
+2
View File
@@ -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)