1 files changed, 4 insertions, 2 deletions
diff --git a/tensorflow/compiler/aot/runtime.cc b/tensorflow/compiler/aot/runtime.cc
index 5e74079fc1..7606420ded 100644
--- a/tensorflow/compiler/aot/runtime.cc
+++ b/tensorflow/compiler/aot/runtime.cc
@@ -64,7 +64,7 @@ size_t align_to(size_t n, size_t align) {
 size_t aligned_buffer_bytes(const intptr_t* sizes, size_t n) {
   size_t total = 0;
   for (size_t i = 0; i < n; ++i) {
-    if (sizes[i] != -1) {
+    if (sizes[i] > 0) {
       total += align_to(sizes[i], kAlign);
     }
   }
@@ -85,7 +85,9 @@ void* MallocContiguousBuffers(const intptr_t* sizes, size_t n, void** bufs,
   }
   uintptr_t pos = reinterpret_cast<uintptr_t>(contiguous);
   for (size_t i = 0; i < n; ++i) {
-    if (sizes[i] == -1) {
+    if (sizes[i] < 0) {
+      // bufs[i] is either a constant, an entry parameter or a thread local
+      // allocation.
       bufs[i] = nullptr;
     } else {
       bufs[i] = reinterpret_cast<void*>(pos);