simplify the KJT.split function when segment is the original KJT (#3014)

Summary:
Pull Request resolved: #3014

# context
* in KJT.split function, when the segment == len(keys), the returned KJT contains the same data as the original KJT
* however in the function it recreates a new one which introduces extra cost
* this diff remove the redundent KJT creation

# analysis
* when segment == len(keys), start has to be zero so the stride_per_key_per_rank is the original one.
* the following KJT init produces the same KJT as self
```
KeyedJaggedTensor(
    keys=self._keys,
    values=self._values,
    weights=self.weights_or_none(),
    lengths=self._lengths,
    offsets=self._offsets,
    stride=self._stride,
    stride_per_key_per_rank=stride_per_key_per_rank,
    stride_per_key=None,
    length_per_key=self._length_per_key,
    lengths_offset_per_key=None,
    offset_per_key=self._offset_per_key,
    index_per_key=self._index_per_key,
    jt_dict=self._jt_dict,
    inverse_indices=None,
)
```

Reviewed By: iamzainhuda

Differential Revision: D70756397

fbshipit-source-id: e299c67792d4219a2220ef52335dcd9fc4cd100d

Author: TroyGarden

torchrec/sparse/jagged_tensor.py

@@ -2354,24 +2354,7 @@ def split(self, segments: List[int]) -> List["KeyedJaggedTensor"]:
             )
             if segment == len(self._keys):
                 # no torch slicing required
-                split_list.append(
-                    KeyedJaggedTensor(
-                        keys=self._keys,
-                        values=self._values,
-                        weights=self.weights_or_none(),
-                        lengths=self._lengths,
-                        offsets=self._offsets,
-                        stride=self._stride,
-                        stride_per_key_per_rank=stride_per_key_per_rank,
-                        stride_per_key=None,
-                        length_per_key=self._length_per_key,
-                        lengths_offset_per_key=None,
-                        offset_per_key=self._offset_per_key,
-                        index_per_key=self._index_per_key,
-                        jt_dict=self._jt_dict,
-                        inverse_indices=None,
-                    )
-                )
+                split_list.append(self)
             elif segment == 0:
                 empty_int_list: List[int] = torch.jit.annotate(List[int], [])
                 split_list.append(