How to define a layer inside a GRU?

I want to replace the layer inside the GRU with a sub-neural network.
At this moment, I just define a GRU using a for-loop with this sub-neural network to process a sequence data. I am afraid that this will make the program slow…
Is there a way to do so?