Make XLAShardedTensor use default __torch_function__ (#6625)

pytorch · Mar 2, 2024 · 5a113af · 5a113af
1 parent 508aa26
commit 5a113af
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/torch_xla/distributed/spmd/xla_sharded_tensor.py b/torch_xla/distributed/spmd/xla_sharded_tensor.py
@@ -168,3 +168,7 @@ def wrap(elem):
       rs = tree_map(wrap,
                     func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs)))
     return rs
+
+  @classmethod
+  def __torch_function__(cls, func, types, args=(), kwargs=None):
+    return super().__torch_function__(func, types, args, kwargs)