diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index 54aea38..5ea7bd9 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -354,7 +354,7 @@
      */
     @Deprecated
     public void reset(Context context) {
-        initGrid(context, getCurrentGridName(context));
+        initGrid(context, getDefaultGridName(context));
     }
 
     @VisibleForTesting
diff --git a/src/com/android/launcher3/model/GridSizeMigrationUtil.java b/src/com/android/launcher3/model/GridSizeMigrationUtil.java
index 8d2a7f9..62ab4a4 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationUtil.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationUtil.java
@@ -17,6 +17,7 @@
 package com.android.launcher3.model;
 
 import static com.android.launcher3.Flags.enableSmartspaceRemovalToggle;
+import static com.android.launcher3.LauncherPrefs.IS_FIRST_LOAD_AFTER_RESTORE;
 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;
@@ -132,7 +133,8 @@
             return true;
         }
 
-        if (Flags.enableGridMigrationFix()
+        if (LauncherPrefs.get(context).get(IS_FIRST_LOAD_AFTER_RESTORE)
+                && Flags.enableGridMigrationFix()
                 && srcDeviceState.getColumns().equals(destDeviceState.getColumns())
                 && srcDeviceState.getRows() < destDeviceState.getRows()) {
             Log.i("b/360462379", "Grid migration fix entry point.");
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index 21897bf..775d248 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -75,7 +75,6 @@
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.LogConfig;
 
-import java.io.File;
 import java.io.InvalidObjectException;
 import java.util.Arrays;
 import java.util.Collection;
@@ -127,12 +126,12 @@
 
         if (Flags.enableNarrowGridRestore()) {
             String oldPhoneFileName = idp.dbFile;
-            List<String> previousDbs = existingDbs();
+            List<String> previousDbs = existingDbs(context);
             removeOldDBs(context, oldPhoneFileName);
             // The idp before this contains data about the old phone, after this it becomes the idp
             // of the current phone.
             idp.reset(context);
-            trySettingPreviousGidAsCurrent(context, idp, oldPhoneFileName, previousDbs);
+            trySettingPreviousGridAsCurrent(context, idp, oldPhoneFileName, previousDbs);
         } else {
             idp.reinitializeAfterRestore(context);
         }
@@ -143,7 +142,7 @@
      * Try setting the gird used in the previous phone to the new one. If the current device doesn't
      * support the previous grid option it will not be set.
      */
-    private static void trySettingPreviousGidAsCurrent(Context context, InvariantDeviceProfile idp,
+    private static void trySettingPreviousGridAsCurrent(Context context, InvariantDeviceProfile idp,
             String oldPhoneDbFileName, List<String> previousDbs) {
         InvariantDeviceProfile.GridOption oldPhoneGridOption = idp.getGridOptionFromFileName(
                 context, oldPhoneDbFileName);
@@ -166,17 +165,19 @@
     /**
      * Returns a list of paths of the existing launcher dbs.
      */
-    private static List<String> existingDbs() {
+    @VisibleForTesting
+    public static List<String> existingDbs(Context context) {
         // At this point idp.dbFile contains the name of the dbFile from the previous phone
         return LauncherFiles.GRID_DB_FILES.stream()
-                .filter(dbName -> new File(dbName).exists())
+                .filter(dbName -> context.getDatabasePath(dbName).exists())
                 .toList();
     }
 
     /**
      * Only keep the last database used on the previous device.
      */
-    private static void removeOldDBs(Context context, String oldPhoneDbFileName) {
+    @VisibleForTesting
+    public static void removeOldDBs(Context context, String oldPhoneDbFileName) {
         // At this point idp.dbFile contains the name of the dbFile from the previous phone
         LauncherFiles.GRID_DB_FILES.stream()
                 .filter(dbName -> !dbName.equals(oldPhoneDbFileName))