diff --git a/tools/templates/RiotXFeature/root/src/app_package/Fragment.kt.ftl b/tools/templates/RiotXFeature/root/src/app_package/Fragment.kt.ftl
index 53dd3a52dc..df69078d9d 100644
--- a/tools/templates/RiotXFeature/root/src/app_package/Fragment.kt.ftl
+++ b/tools/templates/RiotXFeature/root/src/app_package/Fragment.kt.ftl
@@ -20,8 +20,8 @@ data class ${fragmentArgsClass}() : Parcelable
 
 //TODO: add this fragment into FragmentModule
 class ${fragmentClass} @Inject constructor(
-        val viewModelFactory: ${viewModelClass}.Factory
-) : VectorBaseFragment() {
+        private val viewModelFactory: ${viewModelClass}.Factory
+) : VectorBaseFragment(), ${viewModelClass}.Factory by viewModelFactory {
 
 	<#if createFragmentArgs>
 		private val fragmentArgs: ${fragmentArgsClass} by args()
diff --git a/tools/templates/RiotXFeature/root/src/app_package/ViewModel.kt.ftl b/tools/templates/RiotXFeature/root/src/app_package/ViewModel.kt.ftl
index 1e87714655..f4090b40e6 100644
--- a/tools/templates/RiotXFeature/root/src/app_package/ViewModel.kt.ftl
+++ b/tools/templates/RiotXFeature/root/src/app_package/ViewModel.kt.ftl
@@ -1,5 +1,6 @@
 package ${escapeKotlinIdentifiers(packageName)}
 
+import com.airbnb.mvrx.ActivityViewModelContext
 import com.airbnb.mvrx.FragmentViewModelContext
 import com.airbnb.mvrx.MvRxViewModelFactory
 import com.airbnb.mvrx.ViewModelContext
@@ -18,6 +19,7 @@ class ${viewModelClass} @AssistedInject constructor(@Assisted initialState: ${vi
     <#else>
     : VectorViewModel<${viewStateClass}, ${actionClass}, EmptyViewEvents>(initialState) {
     </#if>
+
     @AssistedInject.Factory
     interface Factory {
         fun create(initialState: ${viewStateClass}): ${viewModelClass}
@@ -27,8 +29,11 @@ class ${viewModelClass} @AssistedInject constructor(@Assisted initialState: ${vi
 
         @JvmStatic
         override fun create(viewModelContext: ViewModelContext, state: ${viewStateClass}): ${viewModelClass}? {
-            val fragment: ${fragmentClass} = (viewModelContext as FragmentViewModelContext).fragment()
-            return fragment.viewModelFactory.create(state)
+            val factory = when (viewModelContext) {
+                is FragmentViewModelContext -> viewModelContext.fragment as? Factory
+                is ActivityViewModelContext -> viewModelContext.activity as? Factory
+            }
+            return factory?.create(state) ?: error("You should let your activity/fragment implements Factory interface")
         }
     }
 
diff --git a/vector/src/main/java/im/vector/riotx/core/platform/VectorViewModel.kt b/vector/src/main/java/im/vector/riotx/core/platform/VectorViewModel.kt
index bda4426c45..e82e8b3856 100644
--- a/vector/src/main/java/im/vector/riotx/core/platform/VectorViewModel.kt
+++ b/vector/src/main/java/im/vector/riotx/core/platform/VectorViewModel.kt
@@ -30,6 +30,10 @@ import io.reactivex.Single
 abstract class VectorViewModel<S : MvRxState, VA : VectorViewModelAction, VE : VectorViewEvents>(initialState: S)
     : BaseMvRxViewModel<S>(initialState, false) {
 
+    interface Factory<S: MvRxState> {
+        fun create(state: S): BaseMvRxViewModel<S>
+    }
+
     // Used to post transient events to the View
     protected val _viewEvents = PublishDataSource<VE>()
     val viewEvents: DataSource<VE> = _viewEvents