unwrap

pytorch · Apr 3, 2024 · c06a0c8 · c06a0c8
1 parent 6d4e730
commit c06a0c8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/experimental/torch_xla2/torch_xla2/functions.py b/experimental/torch_xla2/torch_xla2/functions.py
@@ -29,4 +29,4 @@ def __torch_function__(self, func, types, args=(), kwargs=None) -> torch.Tensor:
     if kwargs:
       warnings.warn(f'kwargs not implemented for {kwargs}')
 
-    return tensor.wrap(jax_func(*args))
+    return tensor.wrap(jax_func(*tensor.unwrap(args)))