diff --git a/torchrec/distributed/train_pipeline/__init__.py b/torchrec/distributed/train_pipeline/__init__.py index e3471db03..89c2fe690 100644 --- a/torchrec/distributed/train_pipeline/__init__.py +++ b/torchrec/distributed/train_pipeline/__init__.py @@ -8,7 +8,7 @@ # pyre-strict -from torchrec.distributed.train_pipeline.train_pipeline import ( # noqa +from torchrec.distributed.train_pipeline.train_pipelines import ( # noqa EvalPipelineSparseDist, # noqa PrefetchTrainPipelineSparseDist, # noqa StagedTrainPipeline, # noqa diff --git a/torchrec/distributed/train_pipeline/tests/test_train_pipeline.py b/torchrec/distributed/train_pipeline/tests/test_train_pipeline.py index 43113bc13..a7d523366 100644 --- a/torchrec/distributed/train_pipeline/tests/test_train_pipeline.py +++ b/torchrec/distributed/train_pipeline/tests/test_train_pipeline.py @@ -45,7 +45,7 @@ from torchrec.distributed.tests.test_fp_embeddingbag_utils import ( create_module_and_freeze, ) -from torchrec.distributed.train_pipeline.train_pipeline import ( +from torchrec.distributed.train_pipeline.train_pipelines import ( EvalPipelineSparseDist, PrefetchTrainPipelineSparseDist, StagedTrainPipeline, diff --git a/torchrec/distributed/train_pipeline/train_pipeline.py b/torchrec/distributed/train_pipeline/train_pipelines.py similarity index 100% rename from torchrec/distributed/train_pipeline/train_pipeline.py rename to torchrec/distributed/train_pipeline/train_pipelines.py