diff --git a/src/com/android/launcher3/testing/TestInformationHandler.java b/src/com/android/launcher3/testing/TestInformationHandler.java
index 516c251..de9d4a6 100644
--- a/src/com/android/launcher3/testing/TestInformationHandler.java
+++ b/src/com/android/launcher3/testing/TestInformationHandler.java
@@ -19,6 +19,7 @@
 
 import android.content.Context;
 import android.os.Bundle;
+import android.os.Debug;
 
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.InvariantDeviceProfile;
@@ -114,13 +115,13 @@
                 break;
             }
 
-            case TestProtocol.REQUEST_ALLOCATED_MEMORY: {
-                final Runtime runtime = Runtime.getRuntime();
-                response.putLong(TestProtocol.TEST_INFO_RESPONSE_FIELD,
-                        runtime.totalMemory() - runtime.freeMemory());
+            case TestProtocol.REQUEST_TOTAL_PSS_KB: {
+                Debug.MemoryInfo mem = new Debug.MemoryInfo();
+                Debug.getMemoryInfo(mem);
+                response.putInt(TestProtocol.TEST_INFO_RESPONSE_FIELD, mem.getTotalPss());
                 break;
             }
         }
         return response;
     }
-}
+}
\ No newline at end of file
diff --git a/src/com/android/launcher3/testing/TestProtocol.java b/src/com/android/launcher3/testing/TestProtocol.java
index d569d40..60a59dd 100644
--- a/src/com/android/launcher3/testing/TestProtocol.java
+++ b/src/com/android/launcher3/testing/TestProtocol.java
@@ -73,7 +73,7 @@
     public static final String REQUEST_APP_LIST_FREEZE_FLAGS = "app-list-freeze-flags";
     public static final String REQUEST_OVERVIEW_LEFT_GESTURE_MARGIN = "overview-left-margin";
     public static final String REQUEST_OVERVIEW_RIGHT_GESTURE_MARGIN = "overview-right-margin";
-    public static final String REQUEST_ALLOCATED_MEMORY = "allocated-memory";
+    public static final String REQUEST_TOTAL_PSS_KB = "total_pss";
 
     public static boolean sDebugTracing = false;
     public static final String REQUEST_ENABLE_DEBUG_TRACING = "enable-debug-tracing";
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 0bd6362..14f2c85 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -957,8 +957,8 @@
         getTestInfo(TestProtocol.REQUEST_DISABLE_DEBUG_TRACING);
     }
 
-    public long getAllocatedMemory() {
-        return getTestInfo(TestProtocol.REQUEST_ALLOCATED_MEMORY).
-                getLong(TestProtocol.TEST_INFO_RESPONSE_FIELD);
+    public int getTotalPssKb() {
+        return getTestInfo(TestProtocol.REQUEST_TOTAL_PSS_KB).
+                getInt(TestProtocol.TEST_INFO_RESPONSE_FIELD);
     }
 }
\ No newline at end of file