adding tensordict into targets to avoid package issue (pytorch#2593)

Summary:
chema

# context
* to avoid test failures like the following: https://www.internalfb.com/diff/D65103519?dst_version_fbid=450279051096758
{F1967377517}

# investigation
* error message: P1687414110
* code search: https://fburl.com/code/r3euunl4

# changes
1. add `tensordict` to EBC/EC/KJT targets files
2. use a temporary `try import tensordict ... catch ...` logic to import the `tensordict` module safely
3. set `tensordict` as extern module in the fbpackager

NOTE: the `try import tensordict...catch...` will be removed in the followup diffs where the `tensordict` module is actually needed.

WARNING: This diff also add "tensordict" into fbpackager's extern list, will need to cautiously confirm the change is valid.

Differential Revision: D66465376

torchrec/distributed/embedding.py

@@ -96,6 +96,14 @@
 except OSError:
     pass
 
+try:
+    from tensordict import TensorDict
+except ImportError:
+
+    class TensorDict:
+        pass
+
+
 logger: logging.Logger = logging.getLogger(__name__)
 
 

torchrec/distributed/embeddingbag.py

@@ -99,6 +99,13 @@
 except OSError:
     pass
 
+try:
+    from tensordict import TensorDict
+except ImportError:
+
+    class TensorDict:
+        pass
+
 
 def _pin_and_move(tensor: torch.Tensor, device: torch.device) -> torch.Tensor:
     return (

torchrec/modules/embedding_modules.py

@@ -21,6 +21,14 @@
 from torchrec.sparse.jagged_tensor import JaggedTensor, KeyedJaggedTensor, KeyedTensor
 
 
+try:
+    from tensordict import TensorDict
+except ImportError:
+
+    class TensorDict:
+        pass
+
+
 @torch.fx.wrap
 def reorder_inverse_indices(
     inverse_indices: Optional[Tuple[List[str], torch.Tensor]],

torchrec/sparse/jagged_tensor.py

@@ -49,9 +49,12 @@
 
 # OSS
 try:
-    pass
+    from tensordict import TensorDict
 except ImportError:
-    pass
+
+    class TensorDict:
+        pass
+
 
 logger: logging.Logger = logging.getLogger()