diff --git a/quickstep/src/com/android/quickstep/util/SplitWithKeyboardShortcutController.java b/quickstep/src/com/android/quickstep/util/SplitWithKeyboardShortcutController.java
index 24d8326..f3fa86a 100644
--- a/quickstep/src/com/android/quickstep/util/SplitWithKeyboardShortcutController.java
+++ b/quickstep/src/com/android/quickstep/util/SplitWithKeyboardShortcutController.java
@@ -27,21 +27,15 @@
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
 import android.app.ActivityManager;
-import android.content.Intent;
 import android.graphics.Rect;
 import android.graphics.RectF;
 import android.os.SystemClock;
-import android.os.UserHandle;
-import android.view.View;
 
 import androidx.annotation.BinderThread;
 
-import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.R;
 import com.android.launcher3.anim.PendingAnimation;
-import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.uioverrides.QuickstepLauncher;
-import com.android.quickstep.OverviewCommandHelper;
 import com.android.quickstep.OverviewComponentObserver;
 import com.android.quickstep.RecentsAnimationCallbacks;
 import com.android.quickstep.RecentsAnimationController;
@@ -143,11 +137,7 @@
                     .updateIconInBackground(
                             Task.from(new Task.TaskKey(runningTaskInfo), runningTaskInfo,
                                     false /* isLocked */),
-                            (task) -> {
-                                if (task.thumbnail != null) {
-                                    floatingTaskView.setIcon(task.thumbnail.thumbnail);
-                                }
-                            });
+                            (task) -> floatingTaskView.setIcon(task.icon));
             floatingTaskView.setAlpha(1);
             floatingTaskView.addStagingAnimation(anim, startingTaskRect, mTempRect,
                     false /* fadeWithThumbnail */, true /* isStagedTask */);
diff --git a/quickstep/src/com/android/quickstep/views/FloatingTaskView.java b/quickstep/src/com/android/quickstep/views/FloatingTaskView.java
index a5652dc..f250b8c 100644
--- a/quickstep/src/com/android/quickstep/views/FloatingTaskView.java
+++ b/quickstep/src/com/android/quickstep/views/FloatingTaskView.java
@@ -11,7 +11,6 @@
 import android.graphics.Paint;
 import android.graphics.Rect;
 import android.graphics.RectF;
-import android.graphics.drawable.BitmapDrawable;
 import android.graphics.drawable.Drawable;
 import android.util.AttributeSet;
 import android.util.FloatProperty;
@@ -213,8 +212,8 @@
         mSplitPlaceholderView.getIconView().setRotation(mOrientationHandler.getDegreesRotated());
     }
 
-    public void setIcon(Bitmap icon) {
-        mSplitPlaceholderView.setIcon(new BitmapDrawable(icon), mSplitHolderSize);
+    public void setIcon(Drawable drawable) {
+        mSplitPlaceholderView.setIcon(drawable, mSplitHolderSize);
     }
 
     protected void initPosition(RectF pos, InsettableFrameLayout.LayoutParams lp) {