Instead of just validating the memory type. Fix them if they are not

compatible.
Change: 120497594

1 files changed, 6 insertions, 4 deletions

diff --git a/tensorflow/core/distributed_runtime/graph_mgr.cc b/tensorflow/core/distributed_runtime/graph_mgr.cc
index dd1e9cb625..9ccc00b4e8 100644
--- a/tensorflow/core/distributed_runtime/graph_mgr.cc
+++ b/tensorflow/core/distributed_runtime/graph_mgr.cc
@@ -169,9 +169,10 @@ Status GraphMgr::InitItem(const string& session, const GraphDef& gdef,
     opseg->AddHold(session);
 
     // Function library runtime.
-    unit->lib = NewFunctionLibraryRuntime(
-        unit->device, runner, def->versions().producer(), item->lib_def,
-        graph_options.optimizer_options());
+    unit->lib =
+        NewFunctionLibraryRuntime(worker_env_->device_mgr, unit->device, runner,
+                                  def->versions().producer(), item->lib_def,
+                                  graph_options.optimizer_options());
 
     // Construct the root executor for the subgraph.
     params.device = unit->device;
@@ -199,7 +200,8 @@ Status GraphMgr::InitItem(const string& session, const GraphDef& gdef,
     };
 
     optimizer.Optimize(lib, params.device, &subgraph);
-    s = ValidateMemoryTypes(DeviceType(unit->device->device_type()), subgraph);
+    s = EnsureMemoryTypes(DeviceType(unit->device->device_type()),
+                          unit->device->name(), subgraph);
     if (!s.ok()) {
       delete subgraph;
       break;