From 565943d8e3913564db360045ef9ec43c37c7e47b Mon Sep 17 00:00:00 2001
From: "Edward Z. Yang" <ezyang@meta.com>
Date: Wed, 23 Jul 2025 13:15:42 -0700
Subject: [PATCH] Update

[ghstack-poisoned]
---
 torch/distributed/tensor/_api.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/torch/distributed/tensor/_api.py b/torch/distributed/tensor/_api.py
index b0ee136c135f..9186f54f94da 100644
--- a/torch/distributed/tensor/_api.py
+++ b/torch/distributed/tensor/_api.py
@@ -30,6 +30,7 @@
     Replicate,
     Shard,
 )
+from torch.utils._python_dispatch import return_and_correct_aliasing
 
 
 __all__ = [
@@ -355,10 +356,15 @@ def __coerce_same_metadata_as_tangent__(self, flatten_spec, expected_type=None):
     # pyre-fixme[3]: Return type must be annotated.
     # pyre-fixme[2]: Parameter must be annotated.
     def __torch_dispatch__(cls, func, types, args=(), kwargs=None):  # type: ignore[override]
-        return DTensor._op_dispatcher.dispatch(
+        return return_and_correct_aliasing(
             func,
             args,
-            kwargs or {},
+            kwargs,
+            DTensor._op_dispatcher.dispatch(
+                func,
+                args,
+                kwargs or {},
+            ),
         )
 
     @staticmethod