How to generate the architectural figure of the network in pytorch when the type of the output of the network is list?

Question

I wonder how to use tochviz to generate network architecture, when the output is a list type? the demo code is as follows: