[Graph] Fix placement issue of stage node when enabling async embedding. (#452)

Author: JackMoriarty

tensorflow/python/training/async_embedding_stage.py

@@ -247,7 +247,7 @@ def _perform_stage(self):
 
             stage_outputs[stage_node.name] = stage_node_outputs
             stage_outputs_consumers[stage_node.name] = stage_node_outputs_consumers
-        with ops.device("/job:worker"):
+        with ops.colocate_with(list(self._stage_nodes.keys())[0]):
             stage_output_result = prefetch.staged(stage_outputs,
                                               num_threads=self._threads_num,
                                               capacity=self._capacity,