I am trying to pass an output of a pretrained classifier model to a bunch of conv2D and linear layers (to get some embeddings eventually), but since the classifier outputs a dim of [4,10], I am having trouble passing this to the conv layers.
So, I have:
   def forward(self, data):
         x = data
         x = self.model_ft(x) # model_ft is the pretrained model; 
         print(x.shape)
        # output of the above is : torch.Size([4, 10])
         x = self.conv1(x) # conv2D defined as: self.conv1 = nn.Conv2d(10, 128, 2)
         x = F.relu(x)
         x = self.pool1(x) # defined as: self.pool1 = nn.MaxPool2d(2)
         x = self.conv2(x) # defined as: self.conv2 = nn.Conv2d(128, 64, 5)

However, when I run this, I get:

RuntimeError: Expected 4-dimensional input for 4-dimensional weight [128, 10, 2, 2], but got 2-dimensional input of size [4, 10] instead

Not sure what it is that the conv2D layer is expecting and I could not really figure this out reading the docs   - any pointers would be great!