Fixing broken unit tests regarding module compatibility (#2651)

Summary:
Pull Request resolved: #2651

Fixing unit test breakage in the task: T211156854

Differential Revision: D67581406

fbshipit-source-id: 7e25585e23f6ec3f4924c8ebdaf82ca1ea0d4294

Author: faran928

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 (

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,