Solved by simaiden in post #2 
              
              
                You can’t add the last conv layer because the kernel size is larger than the input in spatial dim (4x4 input and 5x5 kernel). This happen because you are using the conv layer without padding thus the output is smaller than input. 
Try adjusting padding that preserve the input spatial dim