Warning for to_dict when torch.exporting (#2247)

Summary:
Pull Request resolved: #2247

to_dict for non-strict torch.export is extremely slow and can almost always be avoided. This diff adds a warning to notify the user of this case.

Reviewed By: TroyGarden

Differential Revision: D60293485

fbshipit-source-id: 980f9911de4509d6259865dbc85965e2c66589c6

torchrec/sparse/jagged_tensor.py

@@ -8,6 +8,7 @@
 # pyre-strict
 
 import abc
+import logging
 
 import operator
 
@@ -51,6 +52,8 @@
 except ImportError:
     pass
 
+logger: logging.Logger = logging.getLogger()
+
 
 def _pin_and_move(tensor: torch.Tensor, device: torch.device) -> torch.Tensor:
     if is_torchdynamo_compiling():
@@ -2226,6 +2229,10 @@ def __getitem__(self, key: str) -> JaggedTensor:
             )
 
     def to_dict(self) -> Dict[str, JaggedTensor]:
+        if not torch.jit.is_scripting() and is_non_strict_exporting():
+            logger.warn(
+                "Trying to non-strict torch.export KJT to_dict, which is extremely slow and not recommended!"
+            )
         _jt_dict = _maybe_compute_kjt_to_jt_dict(
             stride=self.stride(),
             stride_per_key=self.stride_per_key(),