Thanks for your long answers @Matthew !
This is how I created the visualized models. A link is generated after running the below codes. Simply click it to open the generated graph. PS: You will need to install pydot3 and graphviz for it to work.
Yes, it makes sense! I also read the original Torch implementation, which illustrates exactly the same implementation.