fix typing (#3419)

mosaicml · Jun 21, 2024 · 16e2862 · 16e2862
1 parent 632eb15
commit 16e2862
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/composer/trainer/_patch_pytorch.py b/composer/trainer/_patch_pytorch.py
@@ -933,7 +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
+        from torch.utils._typing_utils import not_none
+        from torch.distributed.device_mesh import DeviceMesh, _mesh_resources
 
         def create_child_mesh(
             self, parent_mesh: 'DeviceMesh', submesh_dim_names: Tuple[str, ...],