diff --git a/vim/.vim/after/ftplugin/python.vim b/vim/.vim/after/ftplugin/python.vim
index 5f3480b..af0112f 100644
--- a/vim/.vim/after/ftplugin/python.vim
+++ b/vim/.vim/after/ftplugin/python.vim
@@ -38,3 +38,26 @@ let g:ale_fixers = {
 \   '*': ['remove_trailing_lines', 'trim_whitespace'],
 \   'python': ['black', 'add_blank_lines_for_python_control_statements'],
 \}
+
+function! LoadVirtualEnv()
+  if filereadable('./.venv/bin/activate')
+    let g:virtualenv_path = getcwd() . '/.venv/bin/python'
+    let g:python3_host_prog = g:virtualenv_path
+    echo "Loaded virtual environment: " . g:virtualenv_path
+  elseif filereadable('./venv/bin/activate')
+    let g:virtualenv_path = getcwd() . '/venv/bin/python'
+    let g:python3_host_prog = g:virtualenv_path
+    echo "Loaded virtual environment: " . g:virtualenv_path
+  else
+    " If no virtualenv is found, fall back to the default Python
+    " Check if variables exist before unsetting them
+    if exists('g:virtualenv_path')
+      unlet g:virtualenv_path
+    endif
+    if exists('g:python3_host_prog')
+      unlet g:python3_host_prog
+    endif
+  endif
+endfunction
+
+autocmd BufEnter * call LoadVirtualEnv()