[TVM FFI] localize load_inline lifetime in tests

Author: SigureMo

test/legacy_test/test_tvm_ffi.py

@@ -63,13 +63,19 @@ def test_c_dlpack_exchange_api_cpu(self):
         """
 
         mod: Module = tvm_ffi.cpp.load_inline(
-            name='mod', cpp_sources=cpp_source, functions='add_one_cpu'
+            name='mod',
+            cpp_sources=cpp_source,
+            functions='add_one_cpu',
+            keep_module_alive=False,
         )
 
-        x = paddle.full((3,), 1.0, dtype='float32').cpu()
-        y = paddle.zeros((3,), dtype='float32').cpu()
-        mod.add_one_cpu(x, y)
-        np.testing.assert_allclose(y.numpy(), [2.0, 2.0, 2.0])
+        def run_check():
+            x = paddle.full((3,), 1.0, dtype='float32').cpu()
+            y = paddle.zeros((3,), dtype='float32').cpu()
+            mod.add_one_cpu(x, y)
+            np.testing.assert_allclose(y.numpy(), [2.0, 2.0, 2.0])
+
+        run_check()
 
     def test_c_dlpack_exchange_api_gpu(self):
         if not paddle.is_compiled_with_cuda():
@@ -116,12 +122,16 @@ def test_c_dlpack_exchange_api_gpu(self):
             cpp_sources=cpp_sources,
             cuda_sources=cuda_sources,
             functions=['add_one_cuda'],
+            keep_module_alive=False,
         )
 
-        x = paddle.full((3,), 1.0, dtype='float32').cuda()
-        y = paddle.zeros((3,), dtype='float32').cuda()
-        mod.add_one_cuda(x, y)
-        np.testing.assert_allclose(y.numpy(), [2.0, 2.0, 2.0])
+        def run_check():
+            x = paddle.full((3,), 1.0, dtype='float32').cuda()
+            y = paddle.zeros((3,), dtype='float32').cuda()
+            mod.add_one_cuda(x, y)
+            np.testing.assert_allclose(y.numpy(), [2.0, 2.0, 2.0])
+
+        run_check()
 
     def test_c_dlpack_exchange_api_alloc_tensor(self):
         cpp_source = r"""
@@ -141,7 +151,10 @@ def test_c_dlpack_exchange_api_alloc_tensor(self):
             }
         """
         mod: Module = tvm_ffi.cpp.load_inline(
-            name='mod', cpp_sources=cpp_source, functions=['add_one_cpu']
+            name='mod',
+            cpp_sources=cpp_source,
+            functions=['add_one_cpu'],
+            keep_module_alive=False,
         )
 
         def run_check():
@@ -200,21 +213,28 @@ def test_data_type_as_input(self):
             }
         """
         mod: Module = tvm_ffi.cpp.load_inline(
-            name='mod', cpp_sources=cpp_source, functions='check_dtype'
+            name='mod',
+            cpp_sources=cpp_source,
+            functions='check_dtype',
+            keep_module_alive=False,
         )
-        for dtype in [
-            paddle.bool,
-            paddle.uint8,
-            paddle.int16,
-            paddle.int32,
-            paddle.int64,
-            paddle.float32,
-            paddle.float64,
-            paddle.float16,
-            paddle.bfloat16,
-        ]:
-            x = paddle.zeros((10,), dtype=dtype).cpu()
-            mod.check_dtype(x, dtype)
+
+        def run_check():
+            for dtype in [
+                paddle.bool,
+                paddle.uint8,
+                paddle.int16,
+                paddle.int32,
+                paddle.int64,
+                paddle.float32,
+                paddle.float64,
+                paddle.float16,
+                paddle.bfloat16,
+            ]:
+                x = paddle.zeros((10,), dtype=dtype).cpu()
+                mod.check_dtype(x, dtype)
+
+        run_check()
 
 
 class TestDLPackDeviceType(unittest.TestCase):
@@ -260,15 +280,21 @@ def test_dlpack_device_type_as_input(self):
             }
         """
         mod: Module = tvm_ffi.cpp.load_inline(
-            name='mod', cpp_sources=cpp_source, functions='check_device'
+            name='mod',
+            cpp_sources=cpp_source,
+            functions='check_device',
+            keep_module_alive=False,
         )
 
-        x_cpu = paddle.zeros((10,), dtype='float32').cpu()
-        mod.check_device(x_cpu, x_cpu.place)
+        def run_check():
+            x_cpu = paddle.zeros((10,), dtype='float32').cpu()
+            mod.check_device(x_cpu, x_cpu.place)
 
-        if paddle.is_compiled_with_cuda():
-            x_gpu = paddle.zeros((10,), dtype='float32').cuda()
-            mod.check_device(x_gpu, x_gpu.place)
+            if paddle.is_compiled_with_cuda():
+                x_gpu = paddle.zeros((10,), dtype='float32').cuda()
+                mod.check_device(x_gpu, x_gpu.place)
+
+        run_check()
 
 
 if __name__ == '__main__':