Update _patch_pytorch.py (#3402)

mosaicml · Jul 21, 2024 · fffa335 · fffa335
1 parent a1c581d
commit fffa335
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/composer/trainer/_patch_pytorch.py b/composer/trainer/_patch_pytorch.py
@@ -933,6 +933,8 @@ def device_mesh__getitem__(self, mesh_dim_names: Union[str, tuple[str]]) -> 'Dev
             return submesh
 
     else:
+        from torch.distributed.device_mesh import _mesh_resources
+
         def create_child_mesh(
             self, parent_mesh: 'DeviceMesh', submesh_dim_names: Tuple[str, ...],
         ) -> 'DeviceMesh':