Training problem

[LoreImbo]

Hi, very nice work. I'm trying to reproduce the results on a new image dataset but I'm experiencing a blocking problem. The problem's output is as follows:

RuntimeError: Expected 4-dimensional input for 4-dimensional weight [64, 3, 3, 3], but got 5-dimensional input of size [16, 1, 3, 448, 448] instead

I find it on line 360 of the 'main.py' file when the function 'train_one_epoch()' is called. Do you have any idea how I can solve it?

I'm pretty sure that [16] is about batch size, because I tried editing several times; while [448,448] is the image size.

Thank you

[HighwayWu]

Thanks for your interest. The input size should be 4-dimensional tensor with [batch_size, 3, img_size, img_size]. So, just simply remove the additional dimension by using "squeeze(1)".

[LoreImbo]

Thank you so much, it worked!