I think my post is related and partially answers your question. The way to do this currently is to use the modulelist and a bunch of IF statements in the forward, that’s what I do now and it works, but the code could be made cleaner I think. I hope that future versions of pytorch address this better.