fix broken torch imports

JKSenthil · facebook-github-bot · commit 993d52471a4c · 2024-09-16T10:48:07.000-07:00
Reviewed By: diego-urgell

Differential Revision: D62662370

fbshipit-source-id: 7a86367ac692a0b81e4bc5fd81d2b917c7660ee6
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:
 

Original file line number	Diff line number	Diff line change
`@@ -32,7 +32,7 @@`
`32`	`32`	`)`
`33`	`33`
`34`	`34`	`try:`
`35`		`- from torch.distributed._tensor._utils import compute_local_shape_and_global_offset`
	`35`	`+ from torch.distributed.tensor._utils import compute_local_shape_and_global_offset`
`36`	`36`
`37`	`37`	`except ImportError:`
`38`	`38`