Commit fe521b06 authored by Benjamin Vandersmissen's avatar Benjamin Vandersmissen
Browse files

Added Adaptive AvgPool to LeNet, such that the model can be used with different sized inputs.

parent 3c13034b
......@@ -11,7 +11,7 @@ class LeNet5(CustomModel):
self.conv1 = add_pruning_info(nn.Conv2d(self.in_channels, 6, 5, padding=2))
self.pool1 = nn.AvgPool2d(2)
self.conv2 = add_pruning_info(nn.Conv2d(6, 16, 5))
self.pool2 = nn.AvgPool2d(2)
self.pool2 = nn.AdaptiveAvgPool2d((5, 5))
self.conv3 = add_pruning_info(nn.Conv2d(16, 120, 5))
self.fc1 = add_pruning_info(nn.Linear(120, 84))
self.fc2 = nn.Linear(84, self.nr_classes)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment