Attach hook using tvm.register_func directly

This is both cleaner and allows using the `override` function parameter to allow overwriting.

tests/python/relay/test_additional_target_hooks.py

@@ -32,7 +32,7 @@
 )
 
 
-def translate_relay_add_to_tir_subtract():
+def translate_relay_add_to_tir_subtract(relay_func, target):
     """A transform to test Relay -> TIR with"""
     ib = tvm.tir.ir_builder.create()
     A = tvm.tir.decl_buffer(
@@ -67,9 +67,7 @@ def translate_relay_add_to_tir_subtract():
     "check_result", [check_vm_result, check_graph_executor_result, check_aot_executor_result]
 )
 def test_tir_external_generation(check_result):
-    @tvm.register_func("target.test.tir_lowering")
-    def relay_to_tir(expr, target):
-        return translate_relay_add_to_tir_subtract(expr)
+    tvm.register_func("target.test.tir_lowering", translate_relay_add_to_tir_subtract, True)
 
     shape = (8, 8)
     x_data = np.random.randint(255, size=shape).astype("float32")