torchrec
torchrec copied to clipboard
Allow a context manager to be called around apply_jit
Summary: When running torch.jit.script on the various forward functions, you can run into issues if there are any other utilites interacting with the function definitions. As an example, if you have another JIT running, you need to disable it throughout this process.
This commit adds the ability to additionally pass an apply_jit_context context manager wherever apply_jit is currently passed that will be called around the application of the torch jit.
Differential Revision: D73781040
This pull request was exported from Phabricator. Differential Revision: D73781040
This pull request was exported from Phabricator. Differential Revision: D73781040
This pull request was exported from Phabricator. Differential Revision: D73781040
This pull request was exported from Phabricator. Differential Revision: D73781040