Fix flaky Github test (#3315)

Summary:
Pull Request resolved: #3315

`_maybe_compute_stride_kjt` returns `stride = 0` when `len(keys) == 0` or when `stride_per_key_per_rank` has invalid values.

Adding a guard condition `if stride > 0:` before the problematic `lengths.view(-1, stride)` operation in the `dist_init` method. When `stride` is 0 (e.g. indicating empty keys), skip the permutation.

Reviewed By: nipung90

Differential Revision: D80948346

fbshipit-source-id: c881ad4571dbc295ab7517bc9e351872bba3970e

Author: jd7-tr

torchrec/sparse/jagged_tensor.py

@@ -2992,18 +2992,19 @@ def dist_init(
                             values.numel(),
                         )
                     elif single_batch_per_rank:
-                        (
-                            lengths,
-                            values,
-                            weights,
-                        ) = torch.ops.fbgemm.permute_2D_sparse_data(
-                            torch.jit._unwrap_optional(recat),
-                            lengths.view(-1, stride),
-                            values,
-                            weights,
-                            values.numel(),
-                        )
-                        lengths = lengths.view(-1)
+                        if stride > 0:
+                            (
+                                lengths,
+                                values,
+                                weights,
+                            ) = torch.ops.fbgemm.permute_2D_sparse_data(
+                                torch.jit._unwrap_optional(recat),
+                                lengths.view(-1, stride),
+                                values,
+                                weights,
+                                values.numel(),
+                            )
+                            lengths = lengths.view(-1)
                     else:  # variable batch size per rank
                         (
                             lengths,