diff --git a/torchrec/distributed/train_pipeline/__init__.py b/torchrec/distributed/train_pipeline/__init__.py
index e49f7ff98..d7b38d2b0 100644
--- a/torchrec/distributed/train_pipeline/__init__.py
+++ b/torchrec/distributed/train_pipeline/__init__.py
@@ -12,6 +12,7 @@
     EvalPipelineSparseDist,  # noqa
     PrefetchTrainPipelineSparseDist,  # noqa
     StagedTrainPipeline,  # noqa
+    TorchCompileConfig,  # noqa
     TrainPipeline,  # noqa
     TrainPipelineBase,  # noqa
     TrainPipelinePT2,  # noqa