diff --git a/src/com/android/launcher3/model/GridSizeMigrationDBController.java b/src/com/android/launcher3/model/GridSizeMigrationDBController.java
index d256d1b..b291421 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationDBController.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationDBController.java
@@ -17,7 +17,6 @@
 package com.android.launcher3.model;
 
 import static com.android.launcher3.Flags.enableSmartspaceRemovalToggle;
-import static com.android.launcher3.Flags.oneGridSpecs;
 import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.LauncherSettings.Favorites.TMP_TABLE;
 import static com.android.launcher3.Utilities.SHOULD_SHOW_FIRST_PAGE_WIDGET;
@@ -39,6 +38,7 @@
 import androidx.annotation.NonNull;
 import androidx.annotation.VisibleForTesting;
 
+import com.android.launcher3.Flags;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.LauncherSettings;
@@ -127,7 +127,7 @@
             return true;
         }
 
-        boolean shouldMigrateToStrictlyTallerGrid = isDestNewDb
+        boolean shouldMigrateToStrictlyTallerGrid = (Flags.oneGridSpecs() || isDestNewDb)
                 && srcDeviceState.getColumns().equals(destDeviceState.getColumns())
                 && srcDeviceState.getRows() < destDeviceState.getRows();
         if (shouldMigrateToStrictlyTallerGrid) {
@@ -142,7 +142,7 @@
             if (shouldMigrateToStrictlyTallerGrid) {
                 // We want to add the extra row(s) to the top of the screen, so we shift the grid
                 // down.
-                if (oneGridSpecs()) {
+                if (Flags.oneGridSpecs()) {
                     shiftWorkspaceByXCells(
                             target.getWritableDatabase(),
                             (destDeviceState.getRows() - srcDeviceState.getRows()),
diff --git a/src/com/android/launcher3/model/GridSizeMigrationLogic.kt b/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
index 876919a..9586bf3 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
+++ b/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
@@ -21,7 +21,6 @@
 import android.util.Log
 import androidx.annotation.VisibleForTesting
 import com.android.launcher3.Flags
-import com.android.launcher3.Flags.oneGridSpecs
 import com.android.launcher3.LauncherPrefs
 import com.android.launcher3.LauncherPrefs.Companion.get
 import com.android.launcher3.LauncherPrefs.Companion.getPrefs
@@ -81,7 +80,7 @@
                 // down.
                 if (shouldMigrateToStrtictlyTallerGrid) {
                     Log.d(TAG, "Migrating to strictly taller grid")
-                    if (oneGridSpecs()) {
+                    if (Flags.oneGridSpecs()) {
                         shiftWorkspaceByXCells(
                             target.writableDatabase,
                             (destDeviceState.rows - srcDeviceState.rows),
@@ -374,7 +373,7 @@
         srcDeviceState: DeviceGridState,
         destDeviceState: DeviceGridState,
     ): Boolean {
-        return isDestNewDb &&
+        return (Flags.oneGridSpecs() || isDestNewDb) &&
             srcDeviceState.columns == destDeviceState.columns &&
             srcDeviceState.rows < destDeviceState.rows
     }