fix device_mesh import

src/llama_cookbook/utils/fsdp_utils.py

@@ -6,7 +6,7 @@ import torch.cuda.nccl as nccl
 import torch.distributed as dist
 
 from llama_cookbook.policies import fpSixteen,bfSixteen, get_llama_wrapper
-from torch.distributed._tensor.device_mesh import init_device_mesh
+from torch.distributed.tensor.device_mesh import init_device_mesh
 
 def fsdp_auto_wrap_policy(model, transformer_layer_names):
     import functools