ValueError: Expected more than 1 value per channel when training, got input size torch.Size([1, 256, 1, 1])

[hasaikeyQAQ]

Dear author, hello,
I encountered an error while running bash train.sh: ValueError: Expected more than 1 value per channel when training, got input size torch. Size ([1, 256, 1, 1]).
The dataset is consistent with yours, and running sh test.sh did not generate any errors.
If you know the solution, please let me know. Thank you very much.