Size mismatch in Conv2d to Linear

I am trying to pass features from a convolution layer to a Fully Connected layer, however I get a size mismatch error although the sizes are correct and fits a convolution transpose layer (i.e. 224 * 224). I get the following error with a mini-batch size set to 1 for debugging.
RuntimeError: size mismatch, m1: [224 x 224], m2: [50176 x 1000] at /opt/conda/conda-bld/pytorch_1524582441669/work/aten/src/THC/generic/THCTensorMathBlas.cu:249

Below is the code:

        self.labelize = nn.Sequential(
            # input is outs x 224 x 224
            nn.Conv2d(outs, 64, kernel_size=1, stride=1),
            nn.BatchNorm2d(64), nn.LeakyReLU(0.2, inplace=True),
            # 64 x 224 x 224
            nn.Conv2d(64, 1, kernel_size=1, stride=1),
            nn.BatchNorm2d(1), nn.LeakyReLU(0.2, inplace=True),
            # state size. 1 x 224 x 224
            nn.Linear(1 * 224 * 224, num_of_classes),
            )

Linear layer expects a 2D tensor, where as conv2d outputs a 4D tensor. Try to use .view() to change 4D tensor to 2D tensor.

1 Like

Thanks… I was assuming it will automatically convert it to 2D tensor