import fbgemm sparse_ops before compilation (pytorch#2144)

Summary:
Pull Request resolved: pytorch#2144

Importing fbgemm.sparse_ops that are needed for torchrec compilation before model.compile call.
The main concern is torch.package tests.

Reviewed By: gnahzg, Microve

Differential Revision: D58830081

fbshipit-source-id: 3497f996c881d9c05814e779e68153b7956dc316

torchrec/distributed/train_pipeline/train_pipelines.py

@@ -243,6 +243,10 @@ def progress(self, dataloader_iter: Iterator[In]) -> Out:
                 torch._dynamo.config.force_unspec_int_unbacked_size_like_on_torchrec_kjt = (
                     True
                 )
+
+                # Importing only before compilation to not slow-done train_pipelines import
+                torch.ops.import_module("fbgemm_gpu.sparse_ops")
+
                 self._model.compile(
                     fullgraph=cc.fullgraph, dynamic=cc.dynamic, backend=cc.backend
                 )