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
import logging
import os
import pickle
import platform
import random
import re
import resource
......@@ -158,6 +159,15 @@ def is_npu() -> bool:
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():
"""
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