Unify invoking JIT functions (#57851)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/57851
The same as the previous PR, but for JIT functions.
ghstack-source-id: 129567069
Test Plan: CI
Reviewed By: mrshenli
Differential Revision: D28253841
fbshipit-source-id: 2b8affde16c106f5c76efa8be49af070213708bf
