From c02eeb969c15a7276e2e6ea1b651d4dff3e41973 Mon Sep 17 00:00:00 2001
From: Weilong Wu <veyron_wu@163.com>
Date: Mon, 4 Apr 2022 15:49:41 +0800
Subject: [PATCH] Updated uva related code (#41391)

---
 paddle/fluid/pybind/eager_method.cc               |  4 ++++
 .../paddle/fluid/dygraph/varbase_patch_methods.py | 15 +++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/paddle/fluid/pybind/eager_method.cc b/paddle/fluid/pybind/eager_method.cc
index 66fba92f67b..1a7eb629a0e 100644
--- a/paddle/fluid/pybind/eager_method.cc
+++ b/paddle/fluid/pybind/eager_method.cc
@@ -1349,6 +1349,10 @@ static PyObject* tensor_method__uva(TensorObject* self, PyObject* args,
                                     PyObject* kwargs) {
   EAGER_TRY
   VLOG(4) << "Running in tensor_method__uva.";
+  PADDLE_ENFORCE_EQ(self->tensor.is_dense_tensor(), true,
+                    platform::errors::InvalidArgument(
+                        "Unified virtual addressing only support "
+                        "DenseTensor currently."));
   PADDLE_ENFORCE_EQ(platform::is_cpu_place(self->tensor.inner_place()), true,
                     platform::errors::InvalidArgument(
                         "Unified virtual addressing only support "
diff --git a/python/paddle/fluid/dygraph/varbase_patch_methods.py b/python/paddle/fluid/dygraph/varbase_patch_methods.py
index c97471d25f1..bd1ca1aa26d 100644
--- a/python/paddle/fluid/dygraph/varbase_patch_methods.py
+++ b/python/paddle/fluid/dygraph/varbase_patch_methods.py
@@ -818,6 +818,21 @@ def monkey_patch_varbase():
 
     @framework.dygraph_only
     def _uva(self, device_id=0):
+        '''
+        Returns self tensor with the UVA(unified virtual addressing).
+
+        Args:
+            device_id(int, optional): The destination GPU device id. Default: None, means current device.
+
+        Examples:
+            .. code-block:: python
+
+              # required: gpu
+              import paddle
+              x = paddle.to_tensor([1, 2, 3], place=paddle.CPUPlace())
+              x._uva()
+              print(x)
+        '''
         self._tensor_uva(device_id)
 
     @framework.dygraph_only
-- 
GitLab