fix unittest

PaddlePaddle · Nov 17, 2020 · 5a6c4a1 · 5a6c4a1 · paddle-bot-old · Nov 18, 2020
1 parent 41d7b76
commit 5a6c4a1
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/python/paddle/vision/models/resnet.py b/python/paddle/vision/models/resnet.py
@@ -245,7 +245,7 @@ def forward(self, x):
         x = self.layer3(x)
         x = self.layer4(x)
 
-        if self.with_pool > 0:
+        if self.with_pool:
             x = self.avgpool(x)
 
         if self.num_classes > 0:

diff --git a/python/paddle/vision/models/vgg.py b/python/paddle/vision/models/vgg.py
@@ -64,7 +64,7 @@ def __init__(self, features, num_classes=1000, with_pool=True):
         if with_pool:
             self.avgpool = nn.AdaptiveAvgPool2D((7, 7))
 
-        if num_classes > 1000:
+        if num_classes > 0:
             self.classifier = nn.Sequential(
                 nn.Linear(512 * 7 * 7, 4096),
                 nn.ReLU(),