Skip to content

Allow a context manager to be called around apply_jit #2927

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

kddnewton
Copy link

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

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Apr 30, 2025
@facebook-github-bot
Copy link
Contributor

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

kddnewton added a commit to kddnewton/torchrec that referenced this pull request May 1, 2025
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 kddnewton force-pushed the export-D73781040 branch from 78f73b6 to 9c67b0c Compare May 1, 2025 13:54
@facebook-github-bot
Copy link
Contributor

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

kddnewton added a commit to kddnewton/torchrec that referenced this pull request May 1, 2025
Summary:
Pull Request resolved: pytorch#2927

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 kddnewton force-pushed the export-D73781040 branch from 9c67b0c to 1cdab0a Compare May 1, 2025 13:57
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.

Reviewed By: SonicField

Differential Revision: D73781040
@kddnewton kddnewton force-pushed the export-D73781040 branch from 1cdab0a to a927a5c Compare May 8, 2025 15:41
@facebook-github-bot
Copy link
Contributor

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. fb-exported
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants