Update squeezenet.py (#172)

changed a small mistake.

Update squeezenet.py (#172)
changed a small mistake.
f4c4d6cb · Sri Krishna · Soumith Chintala · cbb05c59 · f4c4d6cb
Commit f4c4d6cb authored May 21, 2017 by Sri Krishna Committed by Soumith Chintala May 21, 2017
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

torchvision/models/squeezenet.py torchvision/models/squeezenet.py +1 -1

No files found.
--- a/torchvision/models/squeezenet.py
+++ b/torchvision/models/squeezenet.py
@@ -78,7 +78,7 @@ class SqueezeNet(nn.Module):
                Fire(512, 64, 256, 256),
            )
        # Final convolution is initialized differently form the rest
-        final_conv = nn.Conv2d(512, num_classes, kernel_size=1)
+        final_conv = nn.Conv2d(512, self.num_classes, kernel_size=1)
        self.classifier = nn.Sequential(
            nn.Dropout(p=0.5),
            final_conv,