diff --git a/src/com/android/launcher3/dragndrop/AddItemActivity.java b/src/com/android/launcher3/dragndrop/AddItemActivity.java
index 05fdcef461..2873e73d24 100644
--- a/src/com/android/launcher3/dragndrop/AddItemActivity.java
+++ b/src/com/android/launcher3/dragndrop/AddItemActivity.java
@@ -184,7 +184,7 @@ public class AddItemActivity extends BaseActivity
         // user sees
         TextView widgetAppName = findViewById(R.id.widget_appName);
         WidgetSections.WidgetSection section = targetApp.widgetCategory == NO_CATEGORY ? null
-                : WidgetSections.getWidgetSections(this).get(targetApp.widgetCategory);
+                : WidgetSections.get(this).get(targetApp.widgetCategory);
         widgetAppName.setText(section == null ? info.loadLabel(getPackageManager())
                 : getString(section.mSectionTitle));
 
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index 6b3bb5140b..65ff3dc8f2 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -419,7 +419,7 @@ public class LauncherPreviewRenderer extends ContextWrapper
 
     private void inflateAndAddWidgets(LauncherAppWidgetInfo info, WidgetsModel widgetsModel) {
         WidgetItem widgetItem = widgetsModel.getWidgetProviderInfoByProviderName(
-                info.providerName, info.user);
+                info.providerName, info.user, mContext);
         if (widgetItem == null) {
             return;
         }
diff --git a/src/com/android/launcher3/icons/IconCache.java b/src/com/android/launcher3/icons/IconCache.java
index af704a88f0..1862ab24e4 100644
--- a/src/com/android/launcher3/icons/IconCache.java
+++ b/src/com/android/launcher3/icons/IconCache.java
@@ -559,7 +559,7 @@ public class IconCache extends BaseIconCache {
             return;
         }
 
-        WidgetSection widgetSection = WidgetSections.getWidgetSections(mContext)
+        WidgetSection widgetSection = WidgetSections.get(mContext)
                 .get(infoInOut.widgetCategory);
         infoInOut.title = mContext.getString(widgetSection.mSectionTitle);
         infoInOut.contentDescription = getUserBadgedLabel(infoInOut.title, infoInOut.user);
diff --git a/src/com/android/launcher3/model/WidgetsModel.java b/src/com/android/launcher3/model/WidgetsModel.java
index 91ce5eabf5..519eeaa223 100644
--- a/src/com/android/launcher3/model/WidgetsModel.java
+++ b/src/com/android/launcher3/model/WidgetsModel.java
@@ -18,6 +18,7 @@ import android.content.pm.PackageManager;
 import android.os.UserHandle;
 import android.util.Log;
 import android.util.Pair;
+import android.util.SparseArray;
 
 import androidx.annotation.Nullable;
 import androidx.collection.ArrayMap;
@@ -201,8 +202,16 @@ public class WidgetsModel {
         // add and update.
         mWidgetsList.putAll(rawWidgetsShortcuts.stream()
                 .filter(new WidgetValidityCheck(app))
-                .flatMap(widgetItem -> getPackageUserKeys(app.getContext(), widgetItem).stream()
-                        .map(key -> new Pair<>(packageItemInfoCache.getOrCreate(key), widgetItem)))
+                .flatMap(
+                        widgetItem -> getPackageUserKeys(app.getContext(), widgetItem)
+                                .stream()
+                                .map(
+                                        key -> new Pair<>(
+                                                packageItemInfoCache.getOrCreate(key),
+                                                widgetItem
+                                        )
+                                )
+                )
                 .collect(groupingBy(pair -> pair.first, mapping(pair -> pair.second, toList()))));
 
         // Update each package entry
@@ -240,19 +249,26 @@ public class WidgetsModel {
     }
 
     public WidgetItem getWidgetProviderInfoByProviderName(
-            ComponentName providerName, UserHandle user) {
+            ComponentName providerName, UserHandle user, Context context) {
+        SparseArray<WidgetSections.WidgetSection> sections = WidgetSections.get(
+                context);
         if (!WIDGETS_ENABLED) {
             return null;
         }
-        List<WidgetItem> widgetsList = mWidgetsList.get(
-                new PackageItemInfo(providerName.getPackageName(), user));
-        if (widgetsList == null) {
-            return null;
-        }
 
-        for (WidgetItem item : widgetsList) {
-            if (item.componentName.equals(providerName)) {
-                return item;
+        // Checking if we hav ea provider in any of the categories.
+        for (int i = 0; i < sections.size(); i++) {
+            PackageItemInfo key = new PackageItemInfo(
+                    providerName.getPackageName(),
+                    sections.get(i).mCategory,
+                    user
+            );
+            if (mWidgetsList.containsKey(key)) {
+                return mWidgetsList.get(key).stream().filter(
+                                item -> item.componentName.equals(providerName)
+                        )
+                        .findFirst()
+                        .orElse(null);
             }
         }
         return null;
@@ -286,10 +302,12 @@ public class WidgetsModel {
         categories.forEach(category -> {
             if (category == NO_CATEGORY) {
                 packageUserKeys.add(
-                        new PackageUserKey(item.componentName.getPackageName(),
-                                item.user));
+                        new PackageUserKey(item.componentName.getPackageName(), item.user)
+                );
             } else {
-                packageUserKeys.add(new PackageUserKey(category, item.user));
+                packageUserKeys.add(
+                        new PackageUserKey(item.componentName.getPackageName(), category, item.user)
+                );
             }
         });
         return packageUserKeys;
diff --git a/src/com/android/launcher3/util/PackageUserKey.java b/src/com/android/launcher3/util/PackageUserKey.java
index 92d97371e5..c1fb379605 100644
--- a/src/com/android/launcher3/util/PackageUserKey.java
+++ b/src/com/android/launcher3/util/PackageUserKey.java
@@ -48,6 +48,10 @@ public class PackageUserKey {
         update(/* packageName= */ "", widgetCategory, user);
     }
 
+    public PackageUserKey(String packageName, int widgetCategory, UserHandle user) {
+        update(packageName, widgetCategory, user);
+    }
+
     public void update(String packageName, UserHandle user) {
         update(packageName, NO_CATEGORY, user);
     }
diff --git a/src/com/android/launcher3/widget/WidgetSections.java b/src/com/android/launcher3/widget/WidgetSections.java
index c45b09591f..f5fd2a9af9 100644
--- a/src/com/android/launcher3/widget/WidgetSections.java
+++ b/src/com/android/launcher3/widget/WidgetSections.java
@@ -50,7 +50,7 @@ public final class WidgetSections {
     private static Map<ComponentName, IntSet> sWidgetsToCategories;
 
     /** Returns a list of widget sections that are shown in the widget picker. */
-    public static synchronized SparseArray<WidgetSection> getWidgetSections(Context context) {
+    public static synchronized SparseArray<WidgetSection> get(Context context) {
         if (sWidgetSections != null) {
             return sWidgetSections;
         }