diff --git a/torchvision/tv_tensors/_dataset_wrapper.py b/torchvision/tv_tensors/_dataset_wrapper.py index b73cff08b5b..ed178c8d1f9 100644 --- a/torchvision/tv_tensors/_dataset_wrapper.py +++ b/torchvision/tv_tensors/_dataset_wrapper.py @@ -384,12 +384,7 @@ def wrapper(idx, sample): coco_ann = dataset.coco.imgToAnns[image_id] if "masks" in target_keys: target["masks"] = tv_tensors.Mask( - torch.stack( - [ - torch.from_numpy(dataset.coco.annToMask(ann)) - for ann in coco_ann - ] - ), + torch.stack([torch.from_numpy(dataset.coco.annToMask(ann)) for ann in coco_ann]), ) if "labels" in target_keys: