diff --git a/torchrec/distributed/train_pipeline/utils.py b/torchrec/distributed/train_pipeline/utils.py index bb57bbf5e..114f021e2 100644 --- a/torchrec/distributed/train_pipeline/utils.py +++ b/torchrec/distributed/train_pipeline/utils.py @@ -636,6 +636,10 @@ def _rewrite_model( # noqa C901 dist_stream, ) pipelined_forwards.append(child) + else: + logger.warning( + f"Module '{node.target}'' will not be pipelined, due to input modifications" + ) # JIT script unsharded modules if applicable. if apply_jit: