From 777ad774347243709ddcf36e593ac1a0052c3055 Mon Sep 17 00:00:00 2001 From: Sri Krishna Date: Sun, 21 May 2017 19:06:59 +0530 Subject: [PATCH] Update squeezenet.py changed a small mistake. --- torchvision/models/squeezenet.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchvision/models/squeezenet.py b/torchvision/models/squeezenet.py index ddeb7d3150d..3fb48a30082 100644 --- a/torchvision/models/squeezenet.py +++ b/torchvision/models/squeezenet.py @@ -78,7 +78,7 @@ def __init__(self, version=1.0, num_classes=1000): 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,