title: How can I specify the flatten layer input size after many conv layers in PyTorch?
body: <p>Here is my problem, I do a small test on CIFAR10 dataset, how can I specify the flatten layer input size in PyTorch? like the following, the input size is 16*5*5, however I don't know how to calculate this and I want to get the input size through some function.Can someone just write a simple function in this Net class and solve this?</p>
  
  <pre><code>class Net(nn.Module):
      def __init__(self):
          super(Net, self).__init__()
          self.conv1 = nn.Conv2d(3,6,5)  
          self.conv2 = nn.Conv2d(6,16,5)
  
          # HERE , the input size is 16*5*5, but I don't know how to get it.
          self.fc1 = nn.Linear(16*5*5, 120)
          self.fc2 = nn.Linear(120,84)
          self.fc3 = nn.Linear(84,10)
  
      def forward(self, x):
          x = F.max_pool2d(F.relu(self.conv1(x)),(2,2))
          x = F.max_pool2d(F.relu(self.conv2(x)),2)
          x = x.view(x.size()[0],-1)
          x = F.relu(self.fc1(x))
          x = F.relu(self.fc2(x))
          x = self.fc3(x)
          return x
  </code></pre>