diff --git a/test/test_kernel_cache.py b/test/test_kernel_cache.py index b5d30281e5..851f4d83fd 100644 --- a/test/test_kernel_cache.py +++ b/test/test_kernel_cache.py @@ -8,20 +8,22 @@ class TestKernelCache(unittest.TestCase): if Device.DEFAULT not in ["CLANG"]: self.skipTest("No custom kernel cache is implemented") - a = Tensor.rand(4,4) - b = Tensor.rand(4,4) - x = a + b + unique_const = 0.6765677269 + a = Tensor.rand(4,4).realize() + b = Tensor.rand(4,4).realize() + x = a + b + unique_const x.realize() + a1 = Tensor.rand(4,4).realize() + b1 = Tensor.rand(4,4).realize() orig_compile_func = Device['CLANG'].compiler Device['CLANG'].compiler = None # making it not callable - a1 = Tensor.rand(4,4) - b1 = Tensor.rand(4,4) - x1 = a1 + b1 - x1.realize() # Same kernel should be from cache. - - Device['CLANG'].compiler = orig_compile_func + try: + x1 = a1 + b1 + unique_const + x1.realize() # Same kernel should be from cache. + finally: + Device['CLANG'].compiler = orig_compile_func if __name__ == "__main__": unittest.main()