Update paged AdEMAMix

bitsandbytes-foundation · Sep 20, 2024 · 1a3d3c6 · 1a3d3c6
1 parent d4b92d1
commit 1a3d3c6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/bitsandbytes/optim/ademamix.py b/bitsandbytes/optim/ademamix.py
@@ -260,7 +260,7 @@ def update_step(self, group, p, gindex, pindex):
             )
 
     def _get_state_double_buffer(self, p, dtype=torch.float32):
-        if not self.is_paged or p.numel() < 0.5e5:
+        if not self.is_paged or p.numel() < 1e5:
             return torch.zeros((2, *p.size()), dtype=dtype, device=p.device)
         else:
             buff = F.get_paged(*(2, *p.size()), dtype=dtype, device=p.device)