Pt2Pipeline input_transform just before model.compile (#2141)

Summary:
Pull Request resolved: #2141

input_transformer is used to do pt2 hints like torch._dynamo.mark_dynamic/mark_unbacked that should happen exactly on the Tensor objects passed to model.compile

Before that it was called before _to_device() that will produce new Tensor objects.

Reviewed By: TroyGarden

Differential Revision: D58825466

Author: Ivan Kobzarev

torchrec/distributed/train_pipeline/train_pipelines.py

@@ -203,12 +203,14 @@ def progress(self, dataloader_iter: Iterator[In]) -> Out:
         with record_function("## load_batch ##"):
             cur_batch = next(dataloader_iter)
 
-        if self._input_transformer:
-            cur_batch = self._input_transformer(cur_batch)
-
         with record_function("## copy_batch_to_gpu ##"):
             self._cur_batch = _to_device(cur_batch, self._device, non_blocking=False)
 
+        # Input transformer here is used also for pt2 hints to compiler, that should happen on exact object passed to model.compile.
+        # Do not move it before _to_device
+        if self._input_transformer:
+            self._cur_batch = self._input_transformer(self._cur_batch)
+
         if self._model.training:
             with record_function("## zero_grad ##"):
                 self._optimizer.zero_grad()