fix pipeline fwd type (#2962)

xunnanxu · facebook-github-bot · commit 49b55170d168 · 2025-05-08T17:12:40.000-07:00
Summary: Pull Request resolved: #2962 This may lead to ``` assert isinstance(forward, PrefetchPipelinedForward) AssertionError ``` TLDR is the type shouldn't be set at instance level (not sure if this was incorrectly introduced by some kind of rebasing tho...) or inheritance may be accidentally broken. Reviewed By: iamzainhuda Differential Revision: D74428624 fbshipit-source-id: 066029c9978801150c142db7df6cff516aa52047
diff --git a/torchrec/distributed/train_pipeline/train_pipelines.py b/torchrec/distributed/train_pipeline/train_pipelines.py
@@ -468,7 +468,6 @@ def __init__(
         self._model_fwd: Callable[[Optional[In]], Tuple[torch.Tensor, Out]] = (
             custom_model_fwd if custom_model_fwd else model
         )
-        self._pipelined_forward_type = PipelinedForward
         self._dmp_collection_sync_interval_batches = (
             dmp_collection_sync_interval_batches
         )

Original file line number	Diff line number	Diff line change
`@@ -468,7 +468,6 @@ def __init__(`
`468`	`468`	`self._model_fwd: Callable[[Optional[In]], Tuple[torch.Tensor, Out]] = (`
`469`	`469`	`custom_model_fwd if custom_model_fwd else model`
`470`	`470`	`)`
`471`		`- self._pipelined_forward_type = PipelinedForward`
`472`	`471`	`self._dmp_collection_sync_interval_batches = (`
`473`	`472`	`dmp_collection_sync_interval_batches`
`474`	`473`	`)`