diff --git a/neuralnetworks/aidl/vts/functional/Android.bp b/neuralnetworks/aidl/vts/functional/Android.bp
index 7804c2a..d5b150a 100644
--- a/neuralnetworks/aidl/vts/functional/Android.bp
+++ b/neuralnetworks/aidl/vts/functional/Android.bp
@@ -50,9 +50,11 @@
     ],
     static_libs: [
         "android.hardware.common-V2-ndk_platform",
+        "android.hardware.graphics.common-V2-ndk_platform",
         "android.hardware.neuralnetworks-V1-ndk_platform",
         "android.hidl.allocator@1.0",
         "android.hidl.memory@1.0",
+        "libaidlcommonsupport",
         "libgmock",
         "libhidlmemory",
         "libneuralnetworks_generated_test_harness",
diff --git a/neuralnetworks/aidl/vts/functional/MemoryDomainTests.cpp b/neuralnetworks/aidl/vts/functional/MemoryDomainTests.cpp
index 596f8ae..e8313f1 100644
--- a/neuralnetworks/aidl/vts/functional/MemoryDomainTests.cpp
+++ b/neuralnetworks/aidl/vts/functional/MemoryDomainTests.cpp
@@ -16,6 +16,7 @@
 
 #define LOG_TAG "neuralnetworks_aidl_hal_test"
 
+#include <aidl/android/hardware/graphics/common/PixelFormat.h>
 #include <android-base/logging.h>
 #include <android/binder_auto_utils.h>
 #include <android/binder_interface_utils.h>
@@ -659,10 +660,26 @@
         return allocateBuffer(preparedModel, inputIndexes, outputIndexes, {});
     }
 
+    size_t getSize(const Memory& memory) {
+        switch (memory.getTag()) {
+            case Memory::Tag::ashmem:
+                return memory.get<Memory::Tag::ashmem>().size;
+            case Memory::Tag::mappableFile:
+                return memory.get<Memory::Tag::mappableFile>().length;
+            case Memory::Tag::hardwareBuffer: {
+                const auto& hardwareBuffer = memory.get<Memory::Tag::hardwareBuffer>();
+                const bool isBlob =
+                        hardwareBuffer.description.format == graphics::common::PixelFormat::BLOB;
+                return isBlob ? hardwareBuffer.description.width : 0;
+            }
+        }
+        return 0;
+    }
+
     Memory allocateSharedMemory(uint32_t size) {
         const auto sharedMemory = nn::createSharedMemory(size).value();
         auto memory = utils::convert(sharedMemory).value();
-        EXPECT_EQ(memory.size, size);
+        EXPECT_EQ(getSize(memory), size);
         return memory;
     }
 
@@ -690,7 +707,7 @@
 
     void initializeDeviceMemory(const std::shared_ptr<IBuffer>& buffer) {
         Memory memory = allocateSharedMemory(kTestOperandDataSize);
-        ASSERT_EQ(memory.size, kTestOperandDataSize);
+        ASSERT_EQ(getSize(memory), kTestOperandDataSize);
         testCopyFrom(buffer, memory, utils::toSigned(kTestOperand.dimensions).value(),
                      ErrorStatus::NONE);
     }
diff --git a/neuralnetworks/aidl/vts/functional/ValidateModel.cpp b/neuralnetworks/aidl/vts/functional/ValidateModel.cpp
index 94d3daf..698c054 100644
--- a/neuralnetworks/aidl/vts/functional/ValidateModel.cpp
+++ b/neuralnetworks/aidl/vts/functional/ValidateModel.cpp
@@ -259,12 +259,16 @@
 size_t sizeForBinder(const Memory& memory) {
     // This is just a guess.
 
-    size_t size = 0;
-    const NativeHandle& handle = memory.handle;
-    size += sizeof(decltype(handle.fds)::value_type) * handle.fds.size();
-    size += sizeof(decltype(handle.ints)::value_type) * handle.ints.size();
-    size += sizeForBinder(memory.name);
-    size += sizeof(memory);
+    size_t size = sizeof(Memory);
+
+    // Only hardwareBuffer type memory has dynamic memory that needs to be accounted for (in the
+    // form of a NativeHandle type). The other other types of memory (MappableFile, Ashmem) use a
+    // single file descriptor (with metadata) instead.
+    if (memory.getTag() == Memory::Tag::hardwareBuffer) {
+        const NativeHandle& handle = memory.get<Memory::Tag::hardwareBuffer>().handle;
+        size += sizeof(decltype(handle.fds)::value_type) * handle.fds.size();
+        size += sizeof(decltype(handle.ints)::value_type) * handle.ints.size();
+    }
 
     return size;
 }