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,)](