handle empty (keys) sparse features case (#1883)

Summary:
Pull Request resolved: #1883

handle empty keys on meta device path

Reviewed By: MarcioPorto

Differential Revision: D56175026

fbshipit-source-id: 7707cdecc4a68069ee809d630f20aeebbc1db1d7

torchrec/sparse/jagged_tensor.py

@@ -786,7 +786,7 @@ def _maybe_compute_length_per_key(
     values: Optional[torch.Tensor],
 ) -> List[int]:
     if length_per_key is None:
-        if values is not None and values.is_meta:
+        if len(keys) and values is not None and values.is_meta:
             # create dummy lengths per key when on meta device
             total_length = values.numel()
             _length = [total_length // len(keys)] * len(keys)

torchrec/sparse/tests/test_jagged_tensor.py

@@ -1836,6 +1836,13 @@ def test_meta_device_compatibility(self) -> None:
             keys=keys, values=values, weights=weights, offsets=offsets
         )
 
+        # test empty keys case
+        kjt = KeyedJaggedTensor.from_lengths_sync(
+            keys=[],
+            values=torch.tensor([], device=torch.device("meta")),
+            lengths=torch.tensor([], device=torch.device("meta")),
+        )
+
 
 class TestKeyedJaggedTensorScripting(unittest.TestCase):
     def test_scriptable_forward(self) -> None: