Fix OSS test (pytorch#2990)

kausv · facebook-github-bot · commit cef7feaa2a14 · 2025-05-21T14:07:29.000-07:00
Summary: Pull Request resolved: pytorch#2990 Reviewed By: iamzainhuda Differential Revision: D75155654 fbshipit-source-id: bedf34a9b587568ca175d774e2f2a151b05b619c
diff --git a/torchrec/distributed/embedding_lookup.py b/torchrec/distributed/embedding_lookup.py
@@ -14,13 +14,14 @@
 
 import torch
 import torch.distributed as dist
+from fbgemm_gpu.split_table_batched_embeddings_ops_common import BackendType
 from fbgemm_gpu.split_table_batched_embeddings_ops_inference import (
     IntNBitTableBatchedEmbeddingBagsCodegen,
 )
 from fbgemm_gpu.split_table_batched_embeddings_ops_training import (
     SplitTableBatchedEmbeddingBagsCodegen,
 )
-from fbgemm_gpu.tbe.ssd.training import BackendType, SSDTableBatchedEmbeddingBags
+from fbgemm_gpu.tbe.ssd.training import SSDTableBatchedEmbeddingBags
 from fbgemm_gpu.tbe.ssd.utils.partially_materialized_tensor import (
     PartiallyMaterializedTensor,
 )

Original file line number	Diff line number	Diff line change
`@@ -14,13 +14,14 @@`
`14`	`14`
`15`	`15`	`import torch`
`16`	`16`	`import torch.distributed as dist`
	`17`	`+from fbgemm_gpu.split_table_batched_embeddings_ops_common import BackendType`
`17`	`18`	`from fbgemm_gpu.split_table_batched_embeddings_ops_inference import (`
`18`	`19`	`IntNBitTableBatchedEmbeddingBagsCodegen,`
`19`	`20`	`)`
`20`	`21`	`from fbgemm_gpu.split_table_batched_embeddings_ops_training import (`
`21`	`22`	`SplitTableBatchedEmbeddingBagsCodegen,`
`22`	`23`	`)`
`23`		`-from fbgemm_gpu.tbe.ssd.training import BackendType, SSDTableBatchedEmbeddingBags`
	`24`	`+from fbgemm_gpu.tbe.ssd.training import SSDTableBatchedEmbeddingBags`
`24`	`25`	`from fbgemm_gpu.tbe.ssd.utils.partially_materialized_tensor import (`
`25`	`26`	`PartiallyMaterializedTensor,`
`26`	`27`	`)`