Hi,
There is some problem with the input that you feed. In pytorch, we feed input as BxCxHxW. I think in your case, you are feeding channel as the last dimension. The input should be input[64, 32, 32, 3] to have 3 channels, but got 32 channels instead. I think the input should be [64,3,32,32]. You can do this by following line.
X = X.permute((0,3,1,2))

Here X is the input.
Thanks