Remove the stride_per_key_per_rank dim check (#3150)

Summary:
Pull Request resolved: #3150

Fix test failures

Created from CodeHub with https://fburl.com/edit-in-codehub

Reviewed By: spmex

Differential Revision: D77636705

fbshipit-source-id: b44e07dbd592561856fb2cc4ffc360dba3bdcfa2

Author: jd7-tr

torchrec/sparse/jagged_tensor.py

@@ -1767,11 +1767,6 @@ def __init__(
             # does not take List[List[int]]
             assert not isinstance(stride_per_key_per_rank, list)
 
-        if isinstance(stride_per_key_per_rank, torch.IntTensor):
-            assert (
-                stride_per_key_per_rank.dim() == 2
-            ), f"Expect 2D tensor with shape [len(keys), len(ranks)] for stride_per_key_per_rank, but got tensor with shape: {stride_per_key_per_rank.shape}"
-
         self._stride_per_key_per_rank: Optional[torch.IntTensor] = (
             torch.IntTensor(stride_per_key_per_rank, device="cpu")
             if isinstance(stride_per_key_per_rank, list)