Add sanity check for stride_per_key_per_rank's dim (#3124)

Summary:
Pull Request resolved: #3124

`stride_per_key_per_rank` can only be a 2D tensor or 2D list.

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

Reviewed By: TroyGarden

Differential Revision: D77052836

fbshipit-source-id: c06b94c7c8d0999276b5d626feaf434222586860

Author: jd7-tr

torchrec/sparse/jagged_tensor.py

@@ -1761,6 +1761,12 @@ def __init__(
             # in pt2.compile the stride_per_key_per_rank has to be torch.Tensor or None
             # 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)