From eec10aaba2f6c19bf08f3e4983a1dabcca4a5844 Mon Sep 17 00:00:00 2001
From: Zhang Ting <709968123@qq.com>
Date: Tue, 17 Mar 2020 19:27:26 +0800
Subject: [PATCH] set op_device for loss_op_desc (#23027)

---
 python/paddle/fluid/backward.py               |  2 ++
 python/paddle/fluid/optimizer.py              |  3 +--
 .../tests/unittests/test_device_guard.py      | 21 +++++++++++++++++++
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/python/paddle/fluid/backward.py b/python/paddle/fluid/backward.py
index 776af6ceddd..25778e834a2 100644
--- a/python/paddle/fluid/backward.py
+++ b/python/paddle/fluid/backward.py
@@ -295,6 +295,8 @@ def _create_loss_op_desc_(loss):
             core.op_proto_and_checker_maker.kOpRoleAttrName():
             int(core.op_proto_and_checker_maker.OpRole.Backward) |
             int(core.op_proto_and_checker_maker.OpRole.Loss),
+            core.op_proto_and_checker_maker.kOpDeviceAttrName():
+            loss.op.attr(core.op_proto_and_checker_maker.kOpDeviceAttrName())
         })
     return op_desc
 
diff --git a/python/paddle/fluid/optimizer.py b/python/paddle/fluid/optimizer.py
index 2b7275b180f..ff6a42a4235 100644
--- a/python/paddle/fluid/optimizer.py
+++ b/python/paddle/fluid/optimizer.py
@@ -484,8 +484,7 @@ class Optimizer(object):
                     if param_name in input_arg_names:
                         self._param_device_map[param_name] = op.attr(
                             device_attr_name)
-                    else:
-                        self._param_device_map[param_name] = None
+                        break
 
     def _get_device_for_param(self, param_name):
         device = None
diff --git a/python/paddle/fluid/tests/unittests/test_device_guard.py b/python/paddle/fluid/tests/unittests/test_device_guard.py
index 3051ed3995a..1384a22b252 100644
--- a/python/paddle/fluid/tests/unittests/test_device_guard.py
+++ b/python/paddle/fluid/tests/unittests/test_device_guard.py
@@ -144,6 +144,27 @@ class TestDeviceGuard(unittest.TestCase):
         for op in all_ops:
             self.assertEqual(op.desc.attr(device_attr_name), "gpu")
 
+    def test_loss_op_desc(self):
+        main_program = fluid.Program()
+        startup_program = fluid.Program()
+        with fluid.program_guard(main_program, startup_program):
+            data1 = fluid.layers.data(name="data_1", shape=[2], dtype="float32")
+            label = fluid.layers.data(name="label", shape=[1], dtype="int64")
+            fc1 = fluid.layers.fc(input=data1, size=10)
+            with fluid.device_guard("gpu"):
+                out = fluid.layers.softmax_with_cross_entropy(
+                    logits=fc1, label=label)
+                loss = fluid.layers.mean(out)
+                opt = fluid.optimizer.SGDOptimizer(0.1)
+                opt.minimize(loss)
+
+        all_ops = main_program.global_block().ops
+        device_attr_name = core.op_proto_and_checker_maker.kOpDeviceAttrName()
+        for op in all_ops:
+            self.assertEqual(True, op.desc.has_attr(device_attr_name))
+            if op.desc == 'fill_constant':
+                self.assertEqual(op.desc.attr(device_attr_name), "gpu")
+
 
 if __name__ == '__main__':
     unittest.main()
-- 
GitLab