diff options
Diffstat (limited to 'tensorflow/core/framework/types.h')
| -rw-r--r-- | tensorflow/core/framework/types.h | 4 |
1 files changed, 4 insertions, 0 deletions
diff --git a/tensorflow/core/framework/types.h b/tensorflow/core/framework/types.h index ded6aa0991..ff7c9855d6 100644 --- a/tensorflow/core/framework/types.h +++ b/tensorflow/core/framework/types.h @@ -470,6 +470,10 @@ inline bool DataTypeIsUnsigned(DataType dt) { // Returns a 0 on failure int DataTypeSize(DataType dt); +// Returns HOST_MEMORY if `dtype` is always on host or is a DT_INT32, +// DEVICE_MEMORY otherwise. +MemoryType MTypeFromDType(const DataType dtype); + // Types that always sit on host: DT_STRING, DT_STRING_REF, DT_RESOURCE. // For DT_RESOURCE, the handle always sits on host (even if the underlying // object has device-allocated resources). |