This is an interesting question and I’m looking into this recently too.
It depends on how the model was implemented. If the model was implemented in other frameworks (e.g., TensorFlow, PyTorch, etc), then there’s no way for TVM to keep this information, because this hierarchy doesn’t a part of the IR graph but just a nested Python class instance.
If the model was implemented in Relay, then it’s possible to do so by implementing multiple Relay functions. However, I’m not 100% for sure if that would work because some Relay passes may not deal with multiple functions well. More importantly, it may hurt the final end-to-end performance due to unnecessary IR boundaries.