When I run GAN code, I got a runtime error output with shape [1, 28, 28] doesn't match the broadcast shape [3, 28, 28] how fix it