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:39:03 -08:00
commit eddf0c46f4
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

@ -146,8 +146,8 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
num_warps : int
BLOCK_SIZE, num_warps = calculate_settings(n_cols)
Y = torch.empty((n_rows, n_cols), dtype = X.dtype, device = "cuda:0")
r = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0")
Y : torch.Tensor = torch.empty(X.shape, dtype = X.dtype, device = "cuda:0")
r : torch.Tensor = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0")
if not gemma:
_rms_layernorm_forward[(n_rows,)](