What should I do with the flatten layer during back-propagation?
The Flatten layer has no learnable parameters in itself (the operation it performs is fully defined by construction); still, it has to propagate the gradient to the previous layers.
In general, the Flatten operation is well-posed, as whatever is the input shape you know what the output shape is.
When you backpropagate, you are supposed to do an &quot;Unflatten&quot;, which maps a flattened tensor into a tensor of a given shape, and you know what that specific shape is from the forward pass, so it is also a well-posed operation.
More formally
Say you have Img1 in input of your Flatten layer
$$
\begin{pmatrix}
f_{1,1}(x; w_{1,1}) &amp; f_{1,2}(x; w_{1,2}) \\ 
f_{2,1}(x; w_{2,1}) &amp; f_{2,2}(x; w_{2,2})
\end{pmatrix}
$$
So, in the output you have
$$
\begin{pmatrix}
f_{1,1}(x; w_{1,1}) &amp; f_{1,2}(x; w_{1,2}) &amp; f_{2,1}(x; w_{2,1}) &amp; f_{2,2}(x; w_{2,2})
\end{pmatrix}
$$
When you compute the gradient you have
$$
\frac{df_{i,j}(x; w_{i,j})}{dw_{i,j}}
$$
and everything in the same position as in the forward pass, so the unflatten maps from the (1, 4) tensor to the (2, 2) tensor.