set op_device for loss_op_desc (#23027)

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
 

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

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()