Hi Paolo!
Sorry for the delayed reply. Right now we’re on the final stretch before we should have graph mode working end-to-end (see e.g. this stack of PRs: https://github.com/pytorch/pytorch/pull/24426 1), so I don’t see an easy way to offload some of the current work to you right now. This is a critical path, so it can’t be parallelized that much I think. However, once it lands I expect we would discover many places where it can be improved or bugs which need to be fixed - and at that point your help in both finding these spots and helping with fixes would be much appreciated!
Until then I suggest you to get familiar with JIT IR in case you haven’t worked with it much before. A good overview of it can be found here: https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/docs/OVERVIEW.md 12 . I could try to provide more pointers if you need them - let me know if that’s the case.