[Neuron] [Bugfix] Fix neuron startup (#9374)

Co-authored-by: Jerzy Zagorski <jzagorsk@amazon.com>
This commit is contained in:
xendo
2024-10-22 14:51:41 +02:00
committed by GitHub
parent a48e3ec052
commit 9dbcce84a7
7 changed files with 37 additions and 18 deletions

View File

@@ -327,15 +327,6 @@ def is_openvino() -> bool:
return False
@lru_cache(maxsize=None)
def is_neuron() -> bool:
try:
import transformers_neuronx
except ImportError:
transformers_neuronx = None
return transformers_neuronx is not None
@lru_cache(maxsize=None)
def is_xpu() -> bool:
from importlib.metadata import PackageNotFoundError, version
@@ -786,7 +777,7 @@ def is_pin_memory_available() -> bool:
elif is_xpu():
print_warning_once("Pin memory is not supported on XPU.")
return False
elif is_neuron():
elif current_platform.is_neuron():
print_warning_once("Pin memory is not supported on Neuron.")
return False
elif current_platform.is_cpu() or is_openvino():