[CustomDevice] add custom device support for Variable.set_value (#55272)

ronny1996 · web-flow · commit df311526ddd8 · 2023-07-10T17:29:37.000+08:00
diff --git a/paddle/fluid/pybind/place.cc b/paddle/fluid/pybind/place.cc
@@ -640,6 +640,8 @@ void BindPlace(pybind11::module &m) {  // NOLINT
       .def("ipu_device_id", [](platform::Place &self) { return self.device; })
       .def("custom_device_id",
            [](platform::Place &self) { return self.device; })
+      .def("custom_device_type",
+           [](platform::Place &self) { return self.GetDeviceType(); })
       .def("set_place",
            [](platform::Place &self, const platform::Place &other) {
              self = other;
diff --git a/python/paddle/fluid/framework.py b/python/paddle/fluid/framework.py
@@ -2456,6 +2456,12 @@ def set_value(self, value, scope=None):
             p = core.Place()
             p.set_place(t._place())
             place = core.XPUPlace(p.xpu_device_id())
+        elif p.is_custom_place():
+            p = core.Place()
+            p.set_place(t._place())
+            place = core.CustomPlace(
+                p.custom_device_type(), p.custom_device_id()
+            )
         else:
             p = core.Place()
             p.set_place(t._place())
