Shapes cannot be multiplied

I am training my model on ImageNet dataset the backbone network is very big so I decreased the batch size my new error is given below:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (16x25088 and 512x1000)

What changes do I need to make to train my network on ImageNet

class LossNet(nn.Module):
def init(self, feature_sizes=[32, 16, 8, 4], num_channels=[64, 128, 256, 512], interm_dim=128):
super(LossNet, self).init()

    self.GAP1 = nn.AvgPool2d(feature_sizes[0])
    self.GAP2 = nn.AvgPool2d(feature_sizes[1])
    self.GAP3 = nn.AvgPool2d(feature_sizes[2])
    self.GAP4 = nn.AvgPool2d(feature_sizes[3])

    self.FC1 = nn.Linear(num_channels[0], interm_dim)
    self.FC2 = nn.Linear(num_channels[1], interm_dim)
    self.FC3 = nn.Linear(num_channels[2], interm_dim)
    self.FC4 = nn.Linear(num_channels[3], interm_dim)

    self.linear = nn.Linear(4 * interm_dim, 1)

def forward(self, features):
    out1 = self.GAP1(features[0])
    out1 = out1.view(out1.size(0), -1)
    out1 = F.relu(self.FC1(out1))

    out2 = self.GAP2(features[1])
    out2 = out2.view(out2.size(0), -1)
    out2 = F.relu(self.FC2(out2))

    out3 = self.GAP3(features[2])
    out3 = out3.view(out3.size(0), -1)
    out3 = F.relu(self.FC3(out3))

    out4 = self.GAP4(features[3])
    out4 = out4.view(out4.size(0), -1)
    out4 = F.relu(self.FC4(out4))

    out = self.linear(torch.cat((out1, out2, out3, out4), 1))
    return out

This says that a linear layer gets the wrong number of features (and you’d know from the backtrace which). I would propose to print all shapes of interest in the forward to check on it.
Quite likely this is the average pool not doing what you think it does, maybe AdaptiveAveragePool (or just calling “features[0].mean([2, 3])”) is a better fit for your purpose.

Best regards

Thomas