Skip to content

Commit

Permalink
[XPU] API align with new intel pytorch extension release (#6395)
Browse files Browse the repository at this point in the history
with intel_extension_for_pytorch new release, we changed deepspeed
kernel API. Need to align with upstream DeepSpeed xpu op builder.

Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
  • Loading branch information
YizhouZ and loadams authored Aug 22, 2024
1 parent 51da191 commit b81b197
Showing 1 changed file with 5 additions and 2 deletions.
7 changes: 5 additions & 2 deletions op_builder/xpu/inference.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,10 @@ def cxx_args(self):

def load(self):
try:
import intel_extension_for_pytorch.deepspeed
return intel_extension_for_pytorch.deepspeed.transformer_inference.transformer_inference
import intel_extension_for_pytorch
if hasattr(intel_extension_for_pytorch, "deepspeed"):
return intel_extension_for_pytorch.deepspeed.transformer_inference.transformer_inference
else:
return intel_extension_for_pytorch.xpu.deepspeed
except ImportError:
raise ImportError("Please install intel-extension-for-pytorch >= 2.1.30 to include DeepSpeed kernels.")

0 comments on commit b81b197

Please sign in to comment.