2025-01-08 nightly release (65ebab1)

.github/workflows/docs.yml

@@ -58,7 +58,7 @@ jobs:
     - name: Install PyTorch
       shell: bash
       run: |
-        conda install -n build_binary --yes pytorch cpuonly -c pytorch-nightly
+        conda run -n build_binary pip install torch --index-url https://download.pytorch.org/whl/nightly/cpu
     - name: Install fbgemm
       run: |
         conda run -n build_binary pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cpu

.github/workflows/unittest_ci.yml

@@ -63,7 +63,7 @@ jobs:
             python-version: '3.12'
             python-tag: "py312"
             cuda-tag: "cu124"
-    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     with:
       runner: ${{ matrix.os }}
       timeout: 30

.github/workflows/unittest_ci_cpu.yml

@@ -35,7 +35,7 @@ jobs:
             - os: linux.2xlarge
               python-version: '3.12'
               python-tag: "py312"
-    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     with:
       runner: ${{ matrix.os }}
       timeout: 15

torchrec/distributed/embedding.py

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

torchrec/distributed/embeddingbag.py

@@ -102,6 +102,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/distributed/utils.py

@@ -525,7 +525,8 @@ def create_global_tensor_shape_stride_from_metadata(
     """
     size = None
     if parameter_sharding.sharding_type == ShardingType.COLUMN_WISE.value:
-        row_dim = parameter_sharding.sharding_spec.shards[0].shard_sizes[0]  # pyre-ignore[16]
+        # pyre-ignore[16]
+        row_dim = parameter_sharding.sharding_spec.shards[0].shard_sizes[0]
         col_dim = 0
         for shard in parameter_sharding.sharding_spec.shards:
             col_dim += shard.shard_sizes[1]
@@ -551,4 +552,5 @@ def create_global_tensor_shape_stride_from_metadata(
         for _ in range(devices_per_node):
             row_dim += parameter_sharding.sharding_spec.shards[0].shard_sizes[0]
         size = torch.Size([row_dim, col_dim])
-    return size, (size[1], 1) if size else (torch.Size([0, 0]), (0, 1))  # pyre-ignore[7]
+    # pyre-ignore[7]
+    return size, (size[1], 1) if size else (torch.Size([0, 0]), (0, 1))

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()