diff --git a/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
index 91705dc..6282842 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
@@ -43,8 +43,16 @@
         applyButton: Button,
         viewModel: CustomizationOptionsViewModel,
         lifecycleOwner: LifecycleOwner,
+        onNavBack: () -> Unit,
     ) {
-        defaultToolbarBinder.bind(navButton, toolbar, applyButton, viewModel, lifecycleOwner)
+        defaultToolbarBinder.bind(
+            navButton,
+            toolbar,
+            applyButton,
+            viewModel,
+            lifecycleOwner,
+            onNavBack,
+        )
 
         if (viewModel !is ThemePickerCustomizationOptionsViewModel) {
             throw IllegalArgumentException(
@@ -56,7 +64,12 @@
             lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
                 launch {
                     viewModel.onApplyButtonClicked.collect { onApplyButtonClicked ->
-                        applyButton.setOnClickListener { onApplyButtonClicked?.invoke() }
+                        applyButton.setOnClickListener {
+                            if (onApplyButtonClicked != null) {
+                                onApplyButtonClicked.invoke()
+                                onNavBack.invoke()
+                            }
+                        }
                     }
                 }
 
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
index d54fd59..2e500e2 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
@@ -60,13 +60,20 @@
 
     override val selectedOption = defaultCustomizationOptionsViewModel.selectedOption
 
-    override fun deselectOption(): Boolean {
-        keyguardQuickAffordancePickerViewModel2.resetPreview()
-        shapeGridPickerViewModel.resetPreview()
-        clockPickerViewModel.resetPreview()
-        colorPickerViewModel2.resetPreview()
-        darkModeViewModel.resetPreview()
-        return defaultCustomizationOptionsViewModel.deselectOption()
+    override fun handleBackPressed(): Boolean {
+        val isBackPressedHandled = defaultCustomizationOptionsViewModel.handleBackPressed()
+
+        if (isBackPressedHandled) {
+            // If isBackPressedHandled is handled by DefaultCustomizationOptionsViewModel, it means
+            // we navigate back to the main screen from a secondary screen. Reset preview.
+            keyguardQuickAffordancePickerViewModel2.resetPreview()
+            shapeGridPickerViewModel.resetPreview()
+            clockPickerViewModel.resetPreview()
+            colorPickerViewModel2.resetPreview()
+            darkModeViewModel.resetPreview()
+        }
+
+        return isBackPressedHandled
     }
 
     val onCustomizeClockClicked: Flow<(() -> Unit)?> =
@@ -147,16 +154,12 @@
                 }
             }
             .map { onApply ->
-                {
-                    if (onApply != null) {
-                        viewModelScope.launch {
-                            onApply()
-                            // We only wait until onApply() is done to execute deselectOption()
-                            deselectOption()
-                        }
-                    } else {
-                        null
+                if (onApply != null) {
+                    fun() {
+                        viewModelScope.launch { onApply() }
                     }
+                } else {
+                    null
                 }
             }
             .stateIn(viewModelScope, SharingStarted.Eagerly, null)