fix lenet num classes (#28642)

4 years ago · cdc4e6620d
parent 80d2024644
commit cdc4e6620d
1 changed files with 2 additions and 1 deletions
--- a/python/paddle/vision/models/lenet.py
+++ b/python/paddle/vision/models/lenet.py
@ -49,7 +49,8 @@ class LeNet(nn.Layer):

        if num_classes > 0:
            self.fc = nn.Sequential(
-                nn.Linear(400, 120), nn.Linear(120, 84), nn.Linear(84, 10))
+                nn.Linear(400, 120),
+                nn.Linear(120, 84), nn.Linear(84, num_classes))

    def forward(self, inputs):
        x = self.features(inputs)