RuntimeError: Given groups=1, weight of size [64, 3, 3, 3], expected input[1, 512, 512, 3] to have 3 channels, but got 512 channels instead

Experts, is this a problem to convert [1, 512, 512, 3] into [1,3, 512, 512]? If so, where do I need to modify it?

Use (batch_size, channels, height, width) as your input, not (batch_size, height, width, channels).

Thank you for your confirmation. The problem has been solved