diff --git a/lua/layout.lua b/lua/layout.lua
index 2ff572b..434164b 100644
--- a/lua/layout.lua
+++ b/lua/layout.lua
@@ -8,20 +8,18 @@
 local M = {}
 local layouts = {}
 local resize_cmds = {}
-local cursor_positions = {}
 
 function M.save(name)
   layouts[name] = vim.fn.winlayout()
   resize_cmds[name] = vim.fn.winrestcmd()
-  cursor_positions[name] = {}
-  M.add_buf_to_layout(layouts[name], cursor_positions[name])
+  M.add_buf_to_layout(layouts[name])
 end
 
-function M.add_buf_to_layout(layout, cursor_position)
+function M.add_buf_to_layout(layout)
   if layout[1] == "leaf" then
     local win_id = layout[2]
     table.insert(layout, vim.fn.winbufnr(win_id))
-    table.insert(cursor_position, vim.api.nvim_win_get_cursor(win_id))
+    table.insert(layout, vim.api.nvim_win_get_cursor(win_id))
   else
     for _, child_layout in ipairs(layout[2]) do
       M.add_buf_to_layout(child_layout)
@@ -39,21 +37,19 @@ function M.restore(name)
   vim.cmd("wincmd o")
 
   -- Recursively restore buffers and cursor positions
-  M.apply_layout(layouts[name], cursor_positions[name], name)
+  M.apply_layout(layouts[name], name)
 
   -- Resize
   vim.cmd(resize_cmds[name])
 end
 
-function M.apply_layout(layout, cursor_position, name)
+function M.apply_layout(layout, name)
   if layout[1] == "leaf" then
     -- Load buffer for leaf
     if vim.fn.bufexists(layout[3]) == 1 then
       vim.cmd(string.format("b %d", layout[3]))
       -- Restore cursor position
-      if cursor_position then
-        vim.api.nvim_win_set_cursor(vim.fn.win_getid(), cursor_position[1])
-      end
+      vim.api.nvim_win_set_cursor(vim.fn.win_getid(), layout[4])
     end
   else
     -- Split cols or rows, split n-1 times
@@ -73,4 +69,3 @@ function M.apply_layout(layout, cursor_position, name)
 end
 
 return M
-