[fix] match code with ademamix

also define blocksize as in prev commit

Author: cyr0930

bitsandbytes/optim/optimizer.py

@@ -475,9 +475,9 @@ def init_state(self, group, p, gindex, pindex):
             state["qmap2"] = self.name2qmap["udynamic"]
 
             if config["block_wise"]:
+                blocksize = 256
                 n = p.numel()
-                blocks = n // 256
-                blocks += 1 if n % 256 > 0 else 0
+                blocks = (n // blocksize) + bool(n % blocksize)
 
                 state["absmax1"] = torch.zeros((blocks,), dtype=torch.float32, device=p.device)
                 state["absmax2"] = torch.zeros((blocks,), dtype=torch.float32, device=p.device)
@@ -697,9 +697,9 @@ def init_state(self, group, p, gindex, pindex):
             state["qmap1"] = self.name2qmap["dynamic"]
 
             if config["block_wise"]:
+                blocksize = 256
                 n = p.numel()
-                blocks = n // 256
-                blocks += 1 if n % 256 > 0 else 0
+                blocks = (n // blocksize) + bool(n % blocksize)
 
                 state["absmax1"] = torch.zeros((blocks,), dtype=torch.float32, device=p.device)
             else: