remove unecessary wrapper

pytorch · Apr 5, 2024 · 8ea3ec7 · 8ea3ec7
1 parent b08a9ad
commit 8ea3ec7
Showing 1 changed file with 1 addition and 7 deletions.
diff --git a/experimental/torch_xla2/torch_xla2/functions.py b/experimental/torch_xla2/torch_xla2/functions.py
@@ -16,14 +16,8 @@
 def register_function(torch_func: Callable[P, torch.Tensor]):
 
   def decorator(jax_impl: Callable[P, jax.Array]):
-
-    @functools.wraps(torch_func)
-    def wrapper(*args: P.args, **kwargs: P.kwargs):
-      return jax_impl(*args, **kwargs)
-
     registry[torch_func] = jax_impl
-
-    return wrapper
+    return jax_impl
 
   return decorator