Skip to content

Commit 4b02574

Browse files
Lint fix
1 parent 1813b05 commit 4b02574

File tree

3 files changed

+161
-64
lines changed

3 files changed

+161
-64
lines changed

bitsandbytes/backends/default/ops.py

Lines changed: 37 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -320,6 +320,7 @@ def _(
320320
"ademamix": ADEMAMIX,
321321
}
322322

323+
323324
@torch.compile
324325
def _optimizer_precondition_32bit(
325326
g: torch.Tensor,
@@ -525,29 +526,53 @@ def _(
525526

526527
if optimizer_name == "lion":
527528
_optimizer_update_32bit(
528-
g, p, state1, state2, unorm_vec, max_unorm, param_norm,
529-
beta1, beta2, beta3, alpha, eps, weight_decay, step,
530-
lr, gnorm_scale, optimizer_id
529+
g,
530+
p,
531+
state1,
532+
state2,
533+
unorm_vec,
534+
max_unorm,
535+
param_norm,
536+
beta1,
537+
beta2,
538+
beta3,
539+
alpha,
540+
eps,
541+
weight_decay,
542+
step,
543+
lr,
544+
gnorm_scale,
545+
optimizer_id,
531546
)
532547

533548
if max_unorm > 0.0:
534549
unorm_vec.zero_()
535550
_optimizer_precondition_32bit(
536-
g, p, state1, state2, unorm_vec,
537-
beta1, beta2, eps, weight_decay, step,
538-
lr, gnorm_scale, optimizer_id
551+
g, p, state1, state2, unorm_vec, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, optimizer_id
539552
)
540553
else:
541554
if max_unorm > 0.0:
542555
unorm_vec.zero_()
543556
_optimizer_precondition_32bit(
544-
g, p, state1, state2, unorm_vec,
545-
beta1, beta2, eps, weight_decay, step,
546-
lr, gnorm_scale, optimizer_id
557+
g, p, state1, state2, unorm_vec, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, optimizer_id
547558
)
548559

549560
_optimizer_update_32bit(
550-
g, p, state1, state2, unorm_vec, max_unorm, param_norm,
551-
beta1, beta2, beta3, alpha, eps, weight_decay, step,
552-
lr, gnorm_scale, optimizer_id
561+
g,
562+
p,
563+
state1,
564+
state2,
565+
unorm_vec,
566+
max_unorm,
567+
param_norm,
568+
beta1,
569+
beta2,
570+
beta3,
571+
alpha,
572+
eps,
573+
weight_decay,
574+
step,
575+
lr,
576+
gnorm_scale,
577+
optimizer_id,
553578
)

0 commit comments

Comments
 (0)