[misc] add is_cpu() (#6950)
This commit is contained in:
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user