Fixing broken unit tests regarding module compatibility (pytorch#2651)

Summary: Pull Request resolved: pytorch#2651 Fixing unit test breakage in the task: T211156854 Differential Revision: D67581406 fbshipit-source-id: 7e25585e23f6ec3f4924c8ebdaf82ca1ea0d4294
duduyi2013 · Dec 23, 2024 · 5f607ff · 5f607ff
1 parent efca1d6
commit 5f607ff
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/torchrec/distributed/embedding.py b/torchrec/distributed/embedding.py
@@ -28,7 +28,7 @@
 import torch
 from torch import distributed as dist, nn
 from torch.autograd.profiler import record_function
-from torch.distributed._shard.sharding_spec.api import EnumerableShardingSpec
+from torch.distributed._shard.sharding_spec import EnumerableShardingSpec
 from torch.distributed._tensor import DTensor
 from torch.nn.parallel import DistributedDataParallel
 from torchrec.distributed.embedding_sharding import (

diff --git a/torchrec/distributed/quant_embedding.py b/torchrec/distributed/quant_embedding.py
@@ -17,7 +17,7 @@
     IntNBitTableBatchedEmbeddingBagsCodegen,
 )
 from torch import nn
-from torch.distributed._shard.sharding_spec.api import EnumerableShardingSpec
+from torch.distributed._shard.sharding_spec import EnumerableShardingSpec
 from torchrec.distributed.embedding import (
     create_sharding_infos_by_sharding_device_group,
     EmbeddingShardingInfo,