add device check when import ipex

microsoft · Aug 9, 2024 · df9a440 · df9a440
1 parent 61350c5
commit df9a440
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/megatron/model/__init__.py b/megatron/model/__init__.py
@@ -1,8 +1,10 @@
 # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
 
 import torch
-import intel_extension_for_pytorch
 from deepspeed.accelerator.real_accelerator import get_accelerator
+
+if get_accelerator().device_name() == 'xpu':
+    import intel_extension_for_pytorch
 if get_accelerator().device_name() == 'cuda':
     from .fused_layer_norm import MixedFusedLayerNorm as LayerNorm
     from apex.normalization import MixedFusedRMSNorm as RMSNorm