wrap dtensor import for OSS compatibility

Summary: The `compute_local_shape_and_global_offse()` is part of torch nightly, so to prevent any import error for pytorch stable users let's guard this import

Reviewed By: diego-urgell

Differential Revision: D56484333

fbshipit-source-id: 6be460235b7e20ad8dd36091fc4136c091a03397

torchsnapshot/io_preparers/dtensor.py

@@ -30,7 +30,20 @@
     Replicate,
     Shard as ShardPlacement,
 )
-from torch.distributed._tensor._utils import compute_local_shape_and_global_offset
+
+try:
+    from torch.distributed._tensor._utils import compute_local_shape_and_global_offset
+
+except ImportError:
+
+    def compute_local_shape_and_global_offset(
+        global_shape: ShapeType, mesh: DeviceMesh, placements: Sequence[Placement]
+    ) -> Tuple[Tuple[int, ...], Tuple[int, ...]]:
+        raise RuntimeError(
+            "Please use the latest nightly pytorch release to use this feature."
+        )
+
+
 from torchsnapshot.io_preparers.sharded_tensor import (
     _OverlappingRegion,
     ShardedTensorBufferConsumer,