Fix forward compatibility tests & lint error for sharded tensors

Summary: as titled Reviewed By: jarodhou Differential Revision: D34048926 fbshipit-source-id: 383ba4e0b0b2361a9e98fb48b4c876b88860260a
YLGH · Feb 10, 2022 · 6a8f2be · 6a8f2be
1 parent cb6a63c
commit 6a8f2be
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchrec/distributed/types.py b/torchrec/distributed/types.py
@@ -45,15 +45,15 @@ class GenericMeta(type):
 from torch import nn
 
 # @manual
-from torch.distributed._sharded_tensor import (  # noqa
+from torch.distributed._shard.sharded_tensor import (  # noqa
     Shard,
     ShardedTensor,
     ShardedTensorMetadata,
     TensorProperties,
 )
 
 # @manual
-from torch.distributed._sharding_spec import (  # noqa
+from torch.distributed._shard.sharding_spec import (  # noqa
     ShardingSpec,
     ShardMetadata,
     EnumerableShardingSpec,