Update zentorch plugin (#167)

huggingface · Mar 21, 2024 · efa1623 · efa1623
1 parent 38b89e7
commit efa1623
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/optimum_benchmark/backends/pytorch/backend.py b/optimum_benchmark/backends/pytorch/backend.py
@@ -18,7 +18,7 @@
     TrainingArguments,
 )
 
-from ...import_utils import is_deepspeed_available, is_torch_distributed_available, is_torch_zendnn_plugin_available
+from ...import_utils import is_deepspeed_available, is_torch_distributed_available, is_zentorch_available
 from ..base import Backend
 from ..peft_utils import apply_peft
 from ..transformers_utils import random_init_weights
@@ -30,8 +30,8 @@
 if is_torch_distributed_available():
     import torch.distributed
 
-if is_torch_zendnn_plugin_available():
-    import torch_zendnn_plugin  # noqa: F401, F403
+if is_zentorch_available():
+    import zentorch  # type: ignore # noqa: F401
 
 
 # bachend logger

diff --git a/optimum_benchmark/import_utils.py b/optimum_benchmark/import_utils.py
@@ -30,11 +30,11 @@
 _py_txi_available = importlib.util.find_spec("py_txi") is not None
 _pyrsmi_available = importlib.util.find_spec("pyrsmi") is not None
 _llm_swarm_available = importlib.util.find_spec("llm_swarm") is not None
-_torch_zendnn_plugin_available = importlib.util.find_spec("torch_zendnn_plugin") is not None
+_zentorch_available = importlib.util.find_spec("zentorch") is not None
 
 
-def is_torch_zendnn_plugin_available():
-    return _torch_zendnn_plugin_available
+def is_zentorch_available():
+    return _zentorch_available
 
 
 def is_llm_swarm_available():