From db36694eb62d768548434c95bfa4a5289f530a4c Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 4 Nov 2024 23:47:38 -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 c6aecc6d2f..287c8219b5 100644
--- a/unsloth/kernels/rms_layernorm.py
+++ b/unsloth/kernels/rms_layernorm.py
@@ -137,8 +137,8 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
     @staticmethod
     def forward(ctx, X, W, eps : float, gemma : bool = False):
         shape = X.shape
-        dim : int = shape[-1]
-        # X = X.view(-1, dim)
+        # dim : int = shape[-1]
+        X = X.view(shape[0] * shape[1], shape[2])
         n_rows : int
         n_cols : int
         n_rows, n_cols = X.shape