Internal cleanup.

Change: 138135876

1 files changed, 16 insertions, 0 deletions

diff --git a/tensorflow/core/framework/memory_types.cc b/tensorflow/core/framework/memory_types.cc
index 923261c529..14d8d91490 100644
--- a/tensorflow/core/framework/memory_types.cc
+++ b/tensorflow/core/framework/memory_types.cc
@@ -116,6 +116,22 @@ Status MemoryTypesForNode(const OpRegistryInterface* op_registry,
         "' not found in OpDef: ", SummarizeOpDef(*op_def));
   }
 
+  std::vector<int32> hostmem_attr;
+  if (GetNodeAttr(ndef, "_input_hostmem", &hostmem_attr).ok()) {
+    for (int32 i : hostmem_attr) {
+      if (0 <= i && i < inp_mtypes->size()) {
+        (*inp_mtypes)[i] = HOST_MEMORY;
+      }
+    }
+  }
+  if (GetNodeAttr(ndef, "_output_hostmem", &hostmem_attr).ok()) {
+    for (int32 i : hostmem_attr) {
+      if (0 <= i && i < out_mtypes->size()) {
+        (*out_mtypes)[i] = HOST_MEMORY;
+      }
+    }
+  }
+
   return Status::OK();
 }