[MISC] add support custom_op check (#8557)

Co-authored-by: youkaichao <youkaichao@126.com>
This commit is contained in:
Kunshang Ji
2024-09-21 10:03:55 +08:00
committed by GitHub
parent 0057894ef7
commit d4bf085ad0
2 changed files with 33 additions and 22 deletions

View File

@@ -1245,6 +1245,12 @@ def supports_dynamo() -> bool:
return base_torch_version >= Version("2.4.0")
# Some backends use pytorch version < 2.4.0 which doesn't
# support `torch.library.custom_op`.
def supports_custom_op() -> bool:
return hasattr(torch.library, "custom_op")
class AtomicCounter:
"""An atomic, thread-safe counter"""