Sure, I could take a look at your implementation and see where you are stuck.
Generally, you would be able to add layers to each module as long as their setup fits the overall workflow (i.e. in particular the input and output shapes of the new layers should match the model architecture). So feel free to post a code snippet containing the newly added layers, which create the issues.