diff --git a/internal/p10k.zsh b/internal/p10k.zsh
index 4c5206fe..966bcfb9 100755
--- a/internal/p10k.zsh
+++ b/internal/p10k.zsh
@@ -2317,6 +2317,7 @@ powerlevel9k_prepare_prompts() {
 }
 
 function _p9k_zle_keymap_select() {
+  VI_KEYMAP=$KEYMAP
   zle && zle .reset-prompt && zle -R
 }
 
@@ -2826,4 +2827,4 @@ zmodload zsh/mathfunc
 zmodload zsh/system
 zmodload -F zsh/stat b:zstat
 
-prompt_powerlevel9k_setup
\ No newline at end of file
+prompt_powerlevel9k_setup