add lazy getitem to vle and ebc vbe awaitables

Summary: defer wait for lazy awaitable to after result of `__getitem__` is used

Differential Revision: D55765164

torchrec/distributed/embeddingbag.py

@@ -309,7 +309,9 @@ def construct_output_kt(
     )
 
 
-class VariableBatchEmbeddingBagCollectionAwaitable(LazyAwaitable[KeyedTensor]):
+class VariableBatchEmbeddingBagCollectionAwaitable(
+    LazyGetItemMixin[str, torch.Tensor], LazyAwaitable[KeyedTensor]
+):
     def __init__(
         self,
         awaitables: List[Awaitable[torch.Tensor]],