Update rms_layernorm.py

2024-11-04 23:06:13 -08:00 · 2024-11-04 23:06:13 -08:00 · 3fd823bc6c
commit 3fd823bc6c
parent 8a0a732735
1 changed files with 1 additions and 1 deletions
--- 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