ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update fast_lora.py

Browse source
This commit is contained in:
Daniel Han-Chen 2024-01-27 03:38:33 +11:00
commit ba847f5411
1 changed files with 2 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

4
unsloth/kernels/fast_lora.py
View file

@ -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, \