Add num_classes (#128)

1f085a0e · Karan Dwivedi · Soumith Chintala · 74d04d2c · 1f085a0e
Commit 1f085a0e authored Mar 28, 2017 by Karan Dwivedi Committed by Soumith Chintala Mar 27, 2017
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

torchvision/models/vgg.py torchvision/models/vgg.py +2 -2

No files found.
--- a/torchvision/models/vgg.py
+++ b/torchvision/models/vgg.py
@@ -19,7 +19,7 @@ model_urls = {

 class VGG(nn.Module):

-    def __init__(self, features):
+    def __init__(self, features, num_classes=1000):
        super(VGG, self).__init__()
        self.features = features
        self.classifier = nn.Sequential(
@@ -29,7 +29,7 @@ class VGG(nn.Module):
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(),
-            nn.Linear(4096, 1000),
+            nn.Linear(4096, num_classes),
        )
        self._initialize_weights()