To use your code and FlattenerLayer, I had to add
"self.input_space=space"
in the FlattenerLayer set_input_space.

Could you add it and verify there isn't other locations where it is needed?