diff --git a/app/src/main/java/dev/lucasnlm/antimine/custom/CustomLevelDialogFragment.kt b/app/src/main/java/dev/lucasnlm/antimine/custom/CustomLevelDialogFragment.kt
index e59f2f4a..6f3a4488 100644
--- a/app/src/main/java/dev/lucasnlm/antimine/custom/CustomLevelDialogFragment.kt
+++ b/app/src/main/java/dev/lucasnlm/antimine/custom/CustomLevelDialogFragment.kt
@@ -1,8 +1,11 @@
 package dev.lucasnlm.antimine.custom
 
+import android.annotation.SuppressLint
 import android.app.Dialog
 import android.content.DialogInterface
 import android.os.Bundle
+import android.view.LayoutInflater
+import android.view.View
 import android.widget.TextView
 import androidx.appcompat.app.AlertDialog
 import androidx.appcompat.app.AppCompatDialogFragment
@@ -12,44 +15,49 @@ import dev.lucasnlm.antimine.R
 import dev.lucasnlm.antimine.common.level.models.Difficulty
 import dev.lucasnlm.antimine.common.level.models.Minefield
 import dev.lucasnlm.antimine.common.level.viewmodel.GameViewModel
-import dev.lucasnlm.antimine.core.preferences.IPreferencesRepository
 import dev.lucasnlm.antimine.custom.viewmodel.CreateGameViewModel
 import dev.lucasnlm.antimine.custom.viewmodel.CustomEvent
-import javax.inject.Inject
 
 @AndroidEntryPoint
 class CustomLevelDialogFragment : AppCompatDialogFragment() {
-    @Inject
-    lateinit var preferencesRepository: IPreferencesRepository
-
     private val viewModel by activityViewModels<GameViewModel>()
     private val createGameViewModel by activityViewModels<CreateGameViewModel>()
 
-    private fun getSelectedMinefield(): Minefield {
-        val mapWidth: TextView? = dialog?.findViewById(R.id.map_width)
-        val mapHeight: TextView? = dialog?.findViewById(R.id.map_height)
-        val mapMines: TextView? = dialog?.findViewById(R.id.map_mines)
+    private lateinit var mapWidth: TextView
+    private lateinit var mapHeight: TextView
+    private lateinit var mapMines: TextView
 
-        val width = filterInput(
-            mapWidth?.text.toString(),
-            MIN_WIDTH
-        ).coerceAtMost(MAX_WIDTH)
-        val height = filterInput(
-            mapHeight?.text.toString(),
-            MIN_HEIGHT
-        ).coerceAtMost(MAX_HEIGHT)
-        val mines = filterInput(
-            mapMines?.text.toString(),
-            MIN_MINES
-        ).coerceAtMost(width * height - 1)
+    private fun getSelectedMinefield(): Minefield {
+        val width = filterInput(mapWidth.text.toString(), MIN_WIDTH).coerceAtMost(MAX_WIDTH)
+        val height = filterInput(mapHeight.text.toString(), MIN_HEIGHT).coerceAtMost(MAX_HEIGHT)
+        val mines = filterInput(mapMines.text.toString(), MIN_MINES).coerceAtMost(width * height - 1)
 
         return Minefield(width, height, mines)
     }
 
+    @SuppressLint("InflateParams")
+    private fun createView(): View {
+        val view = LayoutInflater
+            .from(context)
+            .inflate(R.layout.dialog_custom_game, null, false)
+
+        mapWidth = view.findViewById(R.id.map_width)
+        mapHeight = view.findViewById(R.id.map_height)
+        mapMines = view.findViewById(R.id.map_mines)
+
+        createGameViewModel.singleState().let {
+            mapWidth.text = it.width.toString()
+            mapHeight.text = it.height.toString()
+            mapMines.text = it.mines.toString()
+        }
+
+        return view
+    }
+
     override fun onCreateDialog(savedInstanceState: Bundle?): Dialog {
         return AlertDialog.Builder(requireContext()).apply {
             setTitle(R.string.new_game)
-            setView(R.layout.dialog_custom_game)
+            setView(createView())
             setNegativeButton(R.string.cancel, null)
             setPositiveButton(R.string.start) { _, _ ->
                 val minefield = getSelectedMinefield()
diff --git a/app/src/main/java/dev/lucasnlm/antimine/custom/viewmodel/CreateGameViewModel.kt b/app/src/main/java/dev/lucasnlm/antimine/custom/viewmodel/CreateGameViewModel.kt
index 871c8adc..6eec91ce 100644
--- a/app/src/main/java/dev/lucasnlm/antimine/custom/viewmodel/CreateGameViewModel.kt
+++ b/app/src/main/java/dev/lucasnlm/antimine/custom/viewmodel/CreateGameViewModel.kt
@@ -2,16 +2,22 @@ package dev.lucasnlm.antimine.custom.viewmodel
 
 import androidx.hilt.lifecycle.ViewModelInject
 import dev.lucasnlm.antimine.core.preferences.IPreferencesRepository
-import dev.lucasnlm.antimine.core.viewmodel.StatelessViewModel
+import dev.lucasnlm.antimine.core.viewmodel.IntentViewModel
+import kotlinx.coroutines.flow.flow
 
 class CreateGameViewModel @ViewModelInject constructor(
     private val preferencesRepository: IPreferencesRepository
-) : StatelessViewModel<CustomEvent>() {
-    override fun onEvent(event: CustomEvent) {
-        when (event) {
-            is CustomEvent.UpdateCustomGameEvent -> {
-                preferencesRepository.updateCustomGameMode(event.minefield)
-            }
+) : IntentViewModel<CustomEvent, CustomState>() {
+
+    override suspend fun mapEventToState(event: CustomEvent) = flow {
+        if (event is CustomEvent.UpdateCustomGameEvent) {
+            val minefield = event.minefield
+            preferencesRepository.updateCustomGameMode(minefield)
+            emit(CustomState(minefield.width, minefield.height, minefield.mines))
         }
     }
+
+    override fun initialState() = with(preferencesRepository.customGameMode()) {
+        CustomState(width, height, mines)
+    }
 }
diff --git a/app/src/main/java/dev/lucasnlm/antimine/custom/viewmodel/CustomState.kt b/app/src/main/java/dev/lucasnlm/antimine/custom/viewmodel/CustomState.kt
new file mode 100644
index 00000000..fe253920
--- /dev/null
+++ b/app/src/main/java/dev/lucasnlm/antimine/custom/viewmodel/CustomState.kt
@@ -0,0 +1,7 @@
+package dev.lucasnlm.antimine.custom.viewmodel
+
+data class CustomState(
+    val width: Int,
+    val height: Int,
+    val mines: Int
+)