Unverified Commit e7759778 authored by JieXin Liang's avatar JieXin Liang Committed by GitHub
Browse files

[misc] add is_cpu() (#6950)

parent 77e928d0
...@@ -25,6 +25,7 @@ import json ...@@ -25,6 +25,7 @@ import json
import logging import logging
import os import os
import pickle import pickle
import platform
import random import random
import re import re
import resource import resource
...@@ -158,6 +159,15 @@ def is_npu() -> bool: ...@@ -158,6 +159,15 @@ def is_npu() -> bool:
return hasattr(torch, "npu") and torch.npu.is_available() return hasattr(torch, "npu") and torch.npu.is_available()
def is_cpu() -> bool:
machine = platform.machine().lower()
return (
machine in ("x86_64", "amd64", "i386", "i686")
and hasattr(torch, "cpu")
and torch.cpu.is_available()
)
def is_flashinfer_available(): def is_flashinfer_available():
""" """
Check whether flashinfer is available. Check whether flashinfer is available.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment