Try this:
class CNN(nn.Module):
def **init**(self):
super(CNN, self).**init**()
self.conv1 = nn.Conv2d(in_channels=3, out_channels=10, kernel_size=3)
self.conv2 = nn.Conv2d(10, 20, kernel_size=3)
self.conv2_drop = nn.Dropout2d()
self.avg_pool=nn.AdaptiveAvgPool2d((6,6)) #adapative pooling
self.fc1 = nn.Linear(720, 1024)
self.fc2 = nn.Linear(1024, 2)
def forward(self, x):
x = F.relu(F.max_pool2d(self.conv1(x), 2))
x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
x=self.avg_pool(x) #apply
x = x.view(x.size(0), -1)
x = F.relu(self.fc1(x))
x = F.dropout(x, training=self.training)
x = self.fc2(x)
return x
Note: The above should only prevent the error and allow some inference.
For reference:
https://pytorch.org/docs/stable/generated/torch.nn.AdaptiveAvgPool2d.html