Merge branch 'master' into models/mobilenetv3

pytorch · Dec 17, 2020 · 5030435 · 5030435
2 parents c189ae1 + 4cbe714
commit 5030435
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchvision/models/mnasnet.py b/torchvision/models/mnasnet.py
@@ -93,11 +93,11 @@ def _get_depths(alpha: float) -> List[int]:
 class MNASNet(torch.nn.Module):
     """ MNASNet, as described in https://arxiv.org/pdf/1807.11626.pdf. This
     implements the B1 variant of the model.
-    >>> model = MNASNet(1000, 1.0)
+    >>> model = MNASNet(1.0, num_classes=1000)
     >>> x = torch.rand(1, 3, 224, 224)
     >>> y = model(x)
     >>> y.dim()
-    1
+    2
     >>> y.nelement()
     1000
     """