From 4fdb4d32ad946d2e1a9c99b726d26228b1696ec0 Mon Sep 17 00:00:00 2001
From: Matthew Douglas <38992547+matthewdouglas@users.noreply.github.com>
Date: Mon, 18 Mar 2024 20:19:09 -0400
Subject: [PATCH] Fix mistake

---
 bitsandbytes/optim/adamw.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/bitsandbytes/optim/adamw.py b/bitsandbytes/optim/adamw.py
index 41314a4f0..f4660588e 100644
--- a/bitsandbytes/optim/adamw.py
+++ b/bitsandbytes/optim/adamw.py
@@ -228,7 +228,7 @@ def step(self, closure=None):
                 self.prefetch_state(p)
 
                 if "rank" in group:
-                    self.update_step(group, p, gindex, pindex, return_updates=lor_update)
+                    self.update_step(group, lor_update, gindex, pindex, return_updates=lor_update)
 
                     # GaLore Projection Back
                     p.data.add_(state["projector"].project_back(lor_update))