From 380aa23ab54c9a9e96c9e34046e2f510c3e18bcc Mon Sep 17 00:00:00 2001 From: Jason Senthil Date: Thu, 25 Apr 2024 11:31:04 -0700 Subject: [PATCH] make torch compile callback OSS compatible Summary: Touches few torch.compile calls in TorchTNT which rely on nightlies to be OSS friendly Reviewed By: galrotem Differential Revision: D56483268 --- torchtnt/framework/callbacks/torch_compile.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/torchtnt/framework/callbacks/torch_compile.py b/torchtnt/framework/callbacks/torch_compile.py index fc5bcf60f3..fb2c744013 100644 --- a/torchtnt/framework/callbacks/torch_compile.py +++ b/torchtnt/framework/callbacks/torch_compile.py @@ -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