diff --git a/sampleapp/src/main/java/com/haroldadmin/sampleapp/AppViewModel.kt b/sampleapp/src/main/java/com/haroldadmin/sampleapp/AppViewModel.kt
index 8790713..81c0606 100644
--- a/sampleapp/src/main/java/com/haroldadmin/sampleapp/AppViewModel.kt
+++ b/sampleapp/src/main/java/com/haroldadmin/sampleapp/AppViewModel.kt
@@ -1,17 +1,14 @@
 package com.haroldadmin.sampleapp
 
-import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.viewModelScope
 import com.haroldadmin.sampleapp.repository.EntitiesRepository
 import com.haroldadmin.vector.VectorState
 import com.haroldadmin.vector.VectorViewModel
-import com.haroldadmin.vector.VectorViewModelFactory
-import com.haroldadmin.vector.ViewModelOwner
 import com.squareup.inject.assisted.Assisted
 import com.squareup.inject.assisted.AssistedInject
 import kotlinx.coroutines.launch
 
-data class AppState(val numberOfEntities: Long) : VectorState
+data class AppState(val numberOfEntities: Long = 0) : VectorState
 
 class AppViewModel @AssistedInject constructor(
     @Assisted initialState: AppState?,
@@ -28,10 +25,4 @@ class AppViewModel @AssistedInject constructor(
     interface Factory {
         fun create(initialState: AppState): AppViewModel
     }
-
-    companion object : VectorViewModelFactory<AppViewModel, AppState> {
-        override fun initialState(handle: SavedStateHandle, owner: ViewModelOwner): AppState? {
-            return AppState(0)
-        }
-    }
 }
\ No newline at end of file
diff --git a/sampleapp/src/main/java/com/haroldadmin/sampleapp/addEditEntity/AddEditEntityFragment.kt b/sampleapp/src/main/java/com/haroldadmin/sampleapp/addEditEntity/AddEditEntityFragment.kt
index 055570c..8249be5 100644
--- a/sampleapp/src/main/java/com/haroldadmin/sampleapp/addEditEntity/AddEditEntityFragment.kt
+++ b/sampleapp/src/main/java/com/haroldadmin/sampleapp/addEditEntity/AddEditEntityFragment.kt
@@ -35,7 +35,6 @@ class AddEditEntityFragment : VectorFragment() {
 
     override fun onAttach(context: Context) {
         AndroidSupportInjection.inject(this)
-//        inject()
         super.onAttach(context)
     }
 
diff --git a/sampleapp/src/main/java/com/haroldadmin/sampleapp/addEditEntity/AddEditEntityViewModel.kt b/sampleapp/src/main/java/com/haroldadmin/sampleapp/addEditEntity/AddEditEntityViewModel.kt
index a9d2d09..395a014 100644
--- a/sampleapp/src/main/java/com/haroldadmin/sampleapp/addEditEntity/AddEditEntityViewModel.kt
+++ b/sampleapp/src/main/java/com/haroldadmin/sampleapp/addEditEntity/AddEditEntityViewModel.kt
@@ -1,5 +1,6 @@
 package com.haroldadmin.sampleapp.addEditEntity
 
+import android.os.Bundle
 import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.viewModelScope
 import com.haroldadmin.sampleapp.CountingEntity
@@ -14,7 +15,7 @@ import com.squareup.inject.assisted.AssistedInject
 import kotlinx.coroutines.launch
 
 class AddEditEntityViewModel @AssistedInject constructor(
-    @Assisted initState: AddEditEntityState? = null,
+    @Assisted initState: AddEditEntityState,
     @Assisted handle: SavedStateHandle,
     private val entityRepository: EntitiesRepository
 ) : SavedStateVectorViewModel<AddEditEntityState>(
@@ -22,18 +23,6 @@ class AddEditEntityViewModel @AssistedInject constructor(
     savedStateHandle = handle
 ) {
 
-    init {
-        if (initState is AddEditEntityState.EditEntity) {
-            viewModelScope.launch {
-                val entity = entityRepository.getEntity(initState.id)
-                setState {
-                    this as AddEditEntityState.EditEntity
-                    copy(name = entity.name, count = entity.counter)
-                }
-            }
-        }
-    }
-
     fun incrementCount() = viewModelScope.launch {
         withState { state ->
             when (state) {
@@ -116,21 +105,42 @@ class AddEditEntityViewModel @AssistedInject constructor(
 
     companion object : VectorViewModelFactory<AddEditEntityViewModel, AddEditEntityState> {
 
+        override fun create(
+            initialState: AddEditEntityState,
+            owner: ViewModelOwner,
+            handle: SavedStateHandle
+        ): AddEditEntityViewModel {
+            throw IllegalStateException("This ViewModel should be created using the AssistedInject Factory only")
+        }
+
         override fun initialState(
             handle: SavedStateHandle,
             owner: ViewModelOwner
         ): AddEditEntityState? {
+
             val persistedState: AddEditEntityState? = handle[KEY_SAVED_STATE]
-            if (persistedState != null) return persistedState
 
-            owner as FragmentViewModelOwner
-            val entityId = owner.args()?.getString("entityId")
+            persistedState?.let {
+                return it
+            } ?: run {
 
-            return if (entityId == null) {
-                AddEditEntityState.AddEntity()
-            } else {
-                AddEditEntityState.EditEntity(id = entityId)
+                val args = (owner as FragmentViewModelOwner).args()
+                val (entityId, entityName, counter) = parseArgs(args!!)
+
+                return if (entityId.isBlank()) {
+                    AddEditEntityState.AddEntity()
+                } else {
+                    AddEditEntityState.EditEntity(entityId, entityName, counter)
+                }
             }
         }
+
+        private fun parseArgs(args: Bundle): Triple<String, String, Long> {
+            val entityId = args.getString("entityId") ?: ""
+            val entityName = args.getString("entityName") ?: ""
+            val count = args.getLong("entityCount")
+
+            return Triple(entityId, entityName, count)
+        }
     }
 }
\ No newline at end of file
diff --git a/sampleapp/src/main/java/com/haroldadmin/sampleapp/entities/EntitiesFragment.kt b/sampleapp/src/main/java/com/haroldadmin/sampleapp/entities/EntitiesFragment.kt
index 40c9e03..e3c5a4b 100644
--- a/sampleapp/src/main/java/com/haroldadmin/sampleapp/entities/EntitiesFragment.kt
+++ b/sampleapp/src/main/java/com/haroldadmin/sampleapp/entities/EntitiesFragment.kt
@@ -34,11 +34,11 @@ class EntitiesFragment : VectorFragment() {
     }
 
     private val entitiesAdapter = EntitiesAdapter(EntitiesDiffCallback()) { entity ->
-        findNavController().navigate(EntitiesFragmentDirections.editEntity(entity.id))
+        findNavController()
+            .navigate(EntitiesFragmentDirections.editEntity(entity.id, entity.name, entity.counter))
     }
 
     override fun onAttach(context: Context) {
-//        inject()
         AndroidSupportInjection.inject(this)
         super.onAttach(context)
     }
diff --git a/sampleapp/src/main/java/com/haroldadmin/sampleapp/entities/EntitiesViewModel.kt b/sampleapp/src/main/java/com/haroldadmin/sampleapp/entities/EntitiesViewModel.kt
index 1a07c0a..21126ce 100644
--- a/sampleapp/src/main/java/com/haroldadmin/sampleapp/entities/EntitiesViewModel.kt
+++ b/sampleapp/src/main/java/com/haroldadmin/sampleapp/entities/EntitiesViewModel.kt
@@ -29,6 +29,15 @@ class EntitiesViewModel @AssistedInject constructor(
     }
 
     companion object : VectorViewModelFactory<EntitiesViewModel, EntitiesState> {
+
+        override fun create(
+            initialState: EntitiesState,
+            owner: ViewModelOwner,
+            handle: SavedStateHandle
+        ): EntitiesViewModel {
+            throw IllegalStateException("This ViewModel should be created using AssistedInject factory only")
+        }
+
         override fun initialState(handle: SavedStateHandle, owner: ViewModelOwner): EntitiesState? {
             return EntitiesState()
         }
diff --git a/sampleapp/src/main/res/navigation/nav_graph.xml b/sampleapp/src/main/res/navigation/nav_graph.xml
index 383da59..1932d27 100644
--- a/sampleapp/src/main/res/navigation/nav_graph.xml
+++ b/sampleapp/src/main/res/navigation/nav_graph.xml
@@ -31,5 +31,15 @@
             app:argType="string"
             app:nullable="true"
             android:defaultValue="@null" />
+        <argument
+            android:name="entityName"
+            app:argType="string"
+            app:nullable="true"
+            android:defaultValue="@null" />
+        <argument
+            android:name="entityCount"
+            app:argType="long"
+            app:nullable="false"
+            android:defaultValue="0L" />
     </fragment>
 </navigation>
\ No newline at end of file