Auto calculating return dimension of squeezenet forward method (#884)

d5347856 · Vitor Finotti Ferreira · Soumith Chintala · 0c575ace · d5347856
Commit d5347856 authored May 03, 2019 by Vitor Finotti Ferreira Committed by Soumith Chintala May 03, 2019
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
@@ -99,7 +99,7 @@ class SqueezeNet(nn.Module):
    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
-        return x.view(x.size(0), self.num_classes)
+        return x.view(x.size(0), -1)


 def _squeezenet(version, pretrained, progress, **kwargs):