RuntimeError: size mismatch, m1: [128 x 2208], m2: [1024 x 512]

Question

I am training an image classification model on the Densenet architecture. My input images are

torch.Size([128, 3, 224, 224])

When I tried training