One more time

bitsandbytes-foundation · Mar 19, 2024 · eceed12 · eceed12
1 parent 4fdb4d3
commit eceed12
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/bitsandbytes/optim/adamw.py b/bitsandbytes/optim/adamw.py
@@ -220,15 +220,15 @@ def step(self, closure=None):
                     lor_update = torch.zeros_like(
                         grad, dtype=p.data.dtype, device=p.data.device, requires_grad=grad.requires_grad
                     )
-                    lor_update.grad = grad
+                    p.grad = grad
 
                 if "state1" not in state:
                     self.init_state(group, p, gindex, pindex)
 
                 self.prefetch_state(p)
 
                 if "rank" in group:
-                    self.update_step(group, lor_update, gindex, pindex, return_updates=lor_update)
+                    self.update_step(group, p, gindex, pindex, return_updates=lor_update)
 
                     # GaLore Projection Back
                     p.data.add_(state["projector"].project_back(lor_update))