add lazy getitem to vle and ebc vbe awaitables

joshuadeng · facebook-github-bot · commit 5e7d32a43b55 · 2024-04-05T12:58:28.000-07:00
Summary: defer wait for lazy awaitable to after result of `__getitem__` is used

Differential Revision: D55765164
diff --git a/torchrec/distributed/embeddingbag.py b/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]],

Original file line number	Diff line number	Diff line change
`@@ -309,7 +309,9 @@ def construct_output_kt(`
`309`	`309`	`)`
`310`	`310`
`311`	`311`
`312`		`-class VariableBatchEmbeddingBagCollectionAwaitable(LazyAwaitable[KeyedTensor]):`
	`312`	`+class VariableBatchEmbeddingBagCollectionAwaitable(`
	`313`	`+ LazyGetItemMixin[str, torch.Tensor], LazyAwaitable[KeyedTensor]`
	`314`	`+):`
`313`	`315`	`def __init__(`
`314`	`316`	`self,`
`315`	`317`	`awaitables: List[Awaitable[torch.Tensor]],`