From 3fd823bc6c786f0902adce0117b3316bc1b9015b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 4 Nov 2024 23:06:13 -0800 Subject: [PATCH] Update rms_layernorm.py --- unsloth/kernels/rms_layernorm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py index 0dc4be6ba7..4074d3a502 100644 --- a/unsloth/kernels/rms_layernorm.py +++ b/unsloth/kernels/rms_layernorm.py @@ -138,7 +138,7 @@ class Fast_RMS_Layernorm(torch.autograd.Function): def forward(ctx, X, W, eps : float, gemma : bool = False): shape = X.shape dim : int = shape[-1] - X = X.view(-1, dim) + X : torch.Tensor = X.view(-1, dim) n_rows : int n_cols : int n_rows, n_cols = X.shape