ソースを参照

fix is_xpu import

Igor Kasianenko 2 ヶ月 前
コミット
b3992f326d
1 ファイル変更2 行追加0 行削除
  1. 2 0
      src/llama_cookbook/utils/fsdp_utils.py

+ 2 - 0
src/llama_cookbook/utils/fsdp_utils.py

@@ -7,6 +7,8 @@ import torch.distributed as dist
 
 from llama_cookbook.policies import fpSixteen,bfSixteen, get_llama_wrapper
 from torch.distributed.device_mesh import init_device_mesh
+from accelerate.utils import is_xpu_available
+
 
 def fsdp_auto_wrap_policy(model, transformer_layer_names):
     import functools