The problem is that your output before applying reshaping is Batchx32x56x56
Why do you reshape that to -1,512?
It makes no sense and that’s why it does not work.
You are squeezing basically maping around 100k elements into a 196*512 matrix