Revert D66465376

Summary: This diff reverts D66465376 seems causing NCCL error and regressing peak mem for a TorchBench test Reviewed By: TroyGarden Differential Revision: D66794877
TroyGarden · Dec 5, 2024 · 2b5db3f · 2b5db3f
1 parent 79111fb
commit 2b5db3f
Show file tree

Hide file tree

Showing 4 changed files with 2 additions and 28 deletions.
diff --git a/torchrec/distributed/embedding.py b/torchrec/distributed/embedding.py
@@ -96,14 +96,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
@@ -99,13 +99,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
@@ -49,12 +49,9 @@
 
 # OSS
 try:
-    from tensordict import TensorDict
+    pass
 except ImportError:
-
-    class TensorDict:
-        pass
-
+    pass
 
 logger: logging.Logger = logging.getLogger()