Skip to content

Commit c5533a2

Browse files
authored
Fix mask dimension in MaskRCNN docstring (#1198)
1 parent 1b673fd commit c5533a2

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

torchvision/models/detection/mask_rcnn.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ class MaskRCNN(FasterRCNN):
3131
- boxes (FloatTensor[N, 4]): the ground-truth boxes in [x1, y1, x2, y2] format, with values
3232
between 0 and H and 0 and W
3333
- labels (Int64Tensor[N]): the class label for each ground-truth box
34-
- masks (UInt8Tensor[N, 1, H, W]): the segmentation binary masks for each instance
34+
- masks (UInt8Tensor[N, H, W]): the segmentation binary masks for each instance
3535
3636
The model returns a Dict[Tensor] during training, containing the classification and regression
3737
losses for both the RPN and the R-CNN, and the mask loss.
@@ -278,7 +278,7 @@ def maskrcnn_resnet50_fpn(pretrained=False, progress=True,
278278
- boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with values
279279
between ``0`` and ``H`` and ``0`` and ``W``
280280
- labels (``Int64Tensor[N]``): the class label for each ground-truth box
281-
- masks (``UInt8Tensor[N, 1, H, W]``): the segmentation binary masks for each instance
281+
- masks (``UInt8Tensor[N, H, W]``): the segmentation binary masks for each instance
282282
283283
The model returns a ``Dict[Tensor]`` during training, containing the classification and regression
284284
losses for both the RPN and the R-CNN, and the mask loss.

0 commit comments

Comments
 (0)