remove temporary import in the try catch

Summary: Tensordict is external dependence from PyTorch atm. this diff is to resolve the failed tests in T212157585 Reviewed By: sarckk Differential Revision: D66772947
TroyGarden · Jan 9, 2025 · 144d79f · 144d79f
1 parent 27e8101
commit 144d79f
Show file tree

Hide file tree

Showing 4 changed files with 0 additions and 30 deletions.
diff --git a/torchrec/distributed/embedding.py b/torchrec/distributed/embedding.py
@@ -97,13 +97,6 @@
 except OSError:
     pass
 
-try:
-    from tensordict import TensorDict
-except ImportError:
-
-    class TensorDict:
-        pass
-
 
 logger: logging.Logger = logging.getLogger(__name__)
 

diff --git a/torchrec/distributed/embeddingbag.py b/torchrec/distributed/embeddingbag.py
@@ -102,13 +102,6 @@
 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 (

diff --git a/torchrec/modules/embedding_modules.py b/torchrec/modules/embedding_modules.py
@@ -21,14 +21,6 @@
 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]],

diff --git a/torchrec/sparse/jagged_tensor.py b/torchrec/sparse/jagged_tensor.py
@@ -47,14 +47,6 @@
 except OSError:
     pass
 
-# OSS
-try:
-    from tensordict import TensorDict
-except ImportError:
-
-    class TensorDict:
-        pass
-
 
 logger: logging.Logger = logging.getLogger()