Skip to content

Commit

Permalink
make torch compile callback OSS compatible
Browse files Browse the repository at this point in the history
Summary: Touches few torch.compile calls in TorchTNT which rely on nightlies to be OSS friendly

Reviewed By: galrotem

Differential Revision: D56483268
  • Loading branch information
JKSenthil authored and facebook-github-bot committed Apr 25, 2024
1 parent 7a3a2f4 commit 380aa23
Showing 1 changed file with 11 additions and 1 deletion.
12 changes: 11 additions & 1 deletion torchtnt/framework/callbacks/torch_compile.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,17 @@

import logging

from torch._inductor.codecache import shutdown_compile_workers
try:
from torch._inductor.codecache import shutdown_compile_workers
except ImportError:

def shutdown_compile_workers() -> None:
logging.warning(
"shutdown_compile_workers is not available in your version of PyTorch. \
Please use nightly version to enable this feature."
)


from torchtnt.framework.callback import Callback
from torchtnt.framework.state import State
from torchtnt.framework.unit import TTrainUnit
Expand Down

0 comments on commit 380aa23

Please sign in to comment.