This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new 3373853 [DLPACK] Skip Loading DLPack if available (#156)
3373853 is described below
commit 33738534976882104799da5ac043ae652685f150
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Oct 16 20:08:52 2025 -0400
[DLPACK] Skip Loading DLPack if available (#156)
This PR updates the optional torch module to
skip loading if attribute exists.
---
python/tvm_ffi/_optional_torch_c_dlpack.py | 11 +++++++++++
1 file changed, 11 insertions(+)
diff --git a/python/tvm_ffi/_optional_torch_c_dlpack.py
b/python/tvm_ffi/_optional_torch_c_dlpack.py
index 3cc1529..6bdff08 100644
--- a/python/tvm_ffi/_optional_torch_c_dlpack.py
+++ b/python/tvm_ffi/_optional_torch_c_dlpack.py
@@ -39,6 +39,17 @@ from . import libinfo
def load_torch_c_dlpack_extension() -> Any:
+ try:
+ import torch # noqa: PLC0415
+
+ if hasattr(torch.Tensor, "__c_dlpack_exchange_api__"):
+ # skip loading the extension if the __c_dlpack_exchange_api__
+ # attribute is already set so we don't have to do it in
+ # newer version of PyTorch
+ return None
+ except ImportError:
+ return None
+
"""Load the torch c dlpack extension."""
cpp_source = """
#include <dlpack/dlpack.h>
