From 9481e2693fbef2b8ae06513e486b7e172df32b30 Mon Sep 17 00:00:00 2001 From: Huanyu He Date: Thu, 6 Feb 2025 11:29:14 -0800 Subject: [PATCH] disable contextlib in PT2 train pipeline (#2730) Summary: Pull Request resolved: https://github.com/pytorch/torchrec/pull/2596 # context * more details in the [post](https://fb.workplace.com/groups/1075192433118967/permalink/1587079018596970/) * disable contextlib with PT2 Differential Revision: D68480678 --- torchrec/distributed/train_pipeline/train_pipelines.py | 1 + 1 file changed, 1 insertion(+) diff --git a/torchrec/distributed/train_pipeline/train_pipelines.py b/torchrec/distributed/train_pipeline/train_pipelines.py index f7b94b37b..2da7c3e86 100644 --- a/torchrec/distributed/train_pipeline/train_pipelines.py +++ b/torchrec/distributed/train_pipeline/train_pipelines.py @@ -269,6 +269,7 @@ def progress(self, dataloader_iter: Iterator[In]) -> Out: # Mandatory dynamo configuration for Torchrec PT2 compilation torch._dynamo.config.capture_scalar_outputs = True + torch._dynamo.config.enable_trace_contextlib = False torch._dynamo.config.capture_dynamic_output_shape_ops = True torch._dynamo.config.force_unspec_int_unbacked_size_like_on_torchrec_kjt = ( True