Sounds like an interesting research direction.
AFAICT your code might be a little more complex than needed. I think you should be able to just use children() recursively, rather than splitting on .. Also, you may even be able to use nn.Module.apply - although I’m not sure if that’ll work in this case.