From ba847f5411b55b948791f251c30b8221ec9727ec Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Sat, 27 Jan 2024 03:38:33 +1100 Subject: [PATCH] Update fast_lora.py --- unsloth/kernels/fast_lora.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index b7649683b3..997235f82e 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -186,7 +186,7 @@ pass class LoRA_MLP_New(torch.autograd.Function): @classmethod @torch.cuda.amp.custom_fwd - def forward(ctx, X : torch.Tensor, + def forward(cls, ctx, X : torch.Tensor, gateW, gateW_quant, gateA, gateB, gateS, upW, upW_quant, upA, upB, upS, downW, downW_quant, downA, downB, downS): @@ -216,7 +216,7 @@ class LoRA_MLP_New(torch.autograd.Function): @classmethod @torch.cuda.amp.custom_bwd - def backward(ctx, dY : torch.Tensor): + def backward(cls, ctx, dY : torch.Tensor): gateW, gateW_quant, gateS, upW, upW_quant, upS, downW, downW_quant, downS, = \ ctx.custom_saved_tensors gateA, gateB, upA, upB, downA, downB, \