fixed function description (#1768)

pytorch · Jan 22, 2020 · 0c92f9c · 0c92f9c
1 parent e2a8b41
commit 0c92f9c
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/torchvision/models/detection/mask_rcnn.py b/torchvision/models/detection/mask_rcnn.py
@@ -222,9 +222,9 @@ class MaskRCNNHeads(nn.Sequential):
     def __init__(self, in_channels, layers, dilation):
         """
         Arguments:
-            num_classes (int): number of output classes
-            input_size (int): number of channels of the input once it's flattened
-            representation_size (int): size of the intermediate representation
+            in_channels (int): number of input channels
+            layers (list): feature dimensions of each FCN layer
+            dilation (int): dilation rate of kernel
         """
         d = OrderedDict()
         next_feature = in_channels