You can’t do that currently in PyTorch. It’s a bit of an administrative problem more than a fundamental one (albeit not an entirely trivial one, I once implemented a PoC for this 18).
Currently, the functional interfaces and computing things in forward is what you have to get along with.
Best regards
Thomas