torchrec icon indicating copy to clipboard operation
torchrec copied to clipboard

Allow a context manager to be called around apply_jit

Open kddnewton opened this issue 7 months ago • 4 comments

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

kddnewton avatar Apr 30 '25 13:04 kddnewton

This pull request was exported from Phabricator. Differential Revision: D73781040

facebook-github-bot avatar Apr 30 '25 13:04 facebook-github-bot

This pull request was exported from Phabricator. Differential Revision: D73781040

facebook-github-bot avatar May 01 '25 13:05 facebook-github-bot

This pull request was exported from Phabricator. Differential Revision: D73781040

facebook-github-bot avatar May 08 '25 15:05 facebook-github-bot

This pull request was exported from Phabricator. Differential Revision: D73781040

facebook-github-bot avatar May 16 '25 16:05 facebook-github-bot