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

code/lenet.py

@@ -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)