diff --git a/tensorflow/lite/micro/simple_memory_allocator.cc b/tensorflow/lite/micro/simple_memory_allocator.cc
index 36ceeafc9d9..d08f48593da 100644
--- a/tensorflow/lite/micro/simple_memory_allocator.cc
+++ b/tensorflow/lite/micro/simple_memory_allocator.cc
@@ -15,6 +15,8 @@ limitations under the License.
 
 #include "tensorflow/lite/micro/simple_memory_allocator.h"
 
+#include <cstddef>
+
 #include "tensorflow/lite/core/api/flatbuffer_conversions.h"
 #include "tensorflow/lite/micro/memory_helpers.h"
 
@@ -29,7 +31,7 @@ uint8_t* SimpleMemoryAllocator::AllocateFromTail(size_t size,
   uint8_t* previous_free = (data_ + data_size_max_) - data_size_;
   uint8_t* current_data = previous_free - size;
   uint8_t* aligned_result = AlignPointerDown(current_data, alignment);
-  size_t aligned_size = (previous_free - aligned_result);
+  std::ptrdiff_t aligned_size = (previous_free - aligned_result);
   if ((data_size_ + aligned_size) > data_size_max_) {
     // TODO(petewarden): Add error reporting beyond returning null!
     return nullptr;
diff --git a/tensorflow/lite/micro/simple_memory_allocator.h b/tensorflow/lite/micro/simple_memory_allocator.h
index 8a4f867c518..c6f0c69fd3f 100644
--- a/tensorflow/lite/micro/simple_memory_allocator.h
+++ b/tensorflow/lite/micro/simple_memory_allocator.h
@@ -35,7 +35,7 @@ class SimpleMemoryAllocator {
   // in ascending order.
   uint8_t* AllocateFromTail(size_t size, size_t alignment);
 
-  int GetDataSize() const { return data_size_; }
+  size_t GetDataSize() const { return data_size_; }
 
   // Child allocator is something like a temporary allocator. Memory allocated
   // by the child allocator will be freed once the child allocator is
@@ -50,7 +50,7 @@ class SimpleMemoryAllocator {
   ~SimpleMemoryAllocator();
 
  private:
-  int data_size_ = 0;
+  size_t data_size_ = 0;
   size_t data_size_max_;
   uint8_t* data_;
   SimpleMemoryAllocator* parent_allocator_ = nullptr;