Update JIT_OPT macro for easier use (#50602)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/50602
Test Plan: Imported from OSS
Reviewed By: ZolotukhinM
Differential Revision: D25931371
fbshipit-source-id: cf6bc58c419a1dc0018639596b304a3a05e38360
