(tvm-ffi) branch main updated: [Fix] Keep a reference to loaded patch lib (#200)

[junrushao]

This is an automated email from the ASF dual-hosted git repository.

junrushao 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 bc2f408  [Fix] Keep a reference to loaded patch lib (#200)
bc2f408 is described below

commit bc2f4088896bc7ca5e356d9cea061e9efca60b54
Author: Yaoyao Ding <dingyaoyao...@gmail.com>
AuthorDate: Wed Oct 29 01:03:23 2025 -0400

    [Fix] Keep a reference to loaded patch lib (#200)
    
    This PRs has a small fix to the previous PR #192 to keep a reference to
    the loaded shared library.
    
    Signed-off-by: Yaoyao Ding <dingyaoyao...@gmail.com>
---
 python/tvm_ffi/_optional_torch_c_dlpack.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/python/tvm_ffi/_optional_torch_c_dlpack.py 
b/python/tvm_ffi/_optional_torch_c_dlpack.py
index c491027..baf100f 100644
--- a/python/tvm_ffi/_optional_torch_c_dlpack.py
+++ b/python/tvm_ffi/_optional_torch_c_dlpack.py
@@ -82,6 +82,8 @@ def load_torch_c_dlpack_extension() -> Any:
 
         # Set the DLPackExchangeAPI pointer on the class
         setattr(torch.Tensor, "__c_dlpack_exchange_api__", func())
+
+        return lib
     except ImportError:
         pass
     except Exception as e:
@@ -110,5 +112,5 @@ def patch_torch_cuda_stream_protocol() -> Any:
 
 
 if os.environ.get("TVM_FFI_DISABLE_TORCH_C_DLPACK", "0") == "0":
-    load_torch_c_dlpack_extension()
+    _LIB = load_torch_c_dlpack_extension()  # keep a reference to the loaded 
shared library
     patch_torch_cuda_stream_protocol()