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

Update rms_layernorm.py

Browse source
This commit is contained in:
Daniel Han 2024-11-04 23:33:08 -08:00
commit 4485a342e0
1 changed files with 2 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

4
unsloth/kernels/rms_layernorm.py
View file

@ -206,9 +206,9 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
pass
def fast_rms_layernorm(layernorm, X, gemma = False):
def fast_rms_layernorm(layernorm, X, gemma : bool = False):
W = layernorm.weight
eps = layernorm.variance_epsilon if \
eps : float = layernorm.variance_epsilon if \
hasattr(layernorm, "variance_epsilon") \
else layernorm.eps
out = Fast_RMS_Layernorm.apply(X, W, eps, gemma)