fix broken torch imports

Reviewed By: diego-urgell Differential Revision: D62662370 fbshipit-source-id: 7a86367ac692a0b81e4bc5fd81d2b917c7660ee6
pytorch · Sep 16, 2024 · 993d524 · 993d524
1 parent d68b3ec
commit 993d524
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/tests/gpu_tests/test_snapshot_dtensor.py b/tests/gpu_tests/test_snapshot_dtensor.py
@@ -13,8 +13,8 @@
 
 import torch
 from torch import distributed as dist, nn
+from torch.distributed import init_device_mesh
 from torch.distributed._tensor import DeviceMesh
-from torch.distributed._tensor.device_mesh import init_device_mesh
 from torch.distributed.fsdp import (
     FullyShardedDataParallel as FSDP,
     ShardingStrategy,

diff --git a/torchsnapshot/io_preparers/dtensor.py b/torchsnapshot/io_preparers/dtensor.py
@@ -32,7 +32,7 @@
 )
 
 try:
-    from torch.distributed._tensor._utils import compute_local_shape_and_global_offset
+    from torch.distributed.tensor._utils import compute_local_shape_and_global_offset
 
 except ImportError: