fix version check bug in get_xpu_available_memory (#3165)

huggingface · Oct 14, 2024 · 292954b · 292954b
1 parent 0e61127
commit 292954b
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/accelerate/utils/memory.py b/src/accelerate/utils/memory.py
@@ -24,6 +24,7 @@
 import warnings
 
 import torch
+from packaging import version
 
 from .imports import (
  is_cuda_available,
@@ -167,7 +168,7 @@ def decorator(*args, **kwargs):
 
 def get_xpu_available_memory(device_index: int):
  if is_ipex_available():
- ipex_version = importlib.metadata.version("intel_extension_for_pytorch")
+ ipex_version = version.parse(importlib.metadata.version("intel_extension_for_pytorch"))
  if compare_versions(ipex_version, ">=", "2.5"):
  from intel_extension_for_pytorch.xpu import mem_get_info