From eddf0c46f409b61ae2053d937fdb5ce08934fc40 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 4 Nov 2024 23:39:03 -0800 Subject: [PATCH] Update rms_layernorm.py --- unsloth/kernels/rms_layernorm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py index c331844e16..ae16b40925 100644 --- a/unsloth/kernels/rms_layernorm.py +++ b/unsloth/kernels/rms_layernorm.py @@ -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,)](