From 4485a342e090928da8e76965ca0f44cc19eb0bb0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 4 Nov 2024 23:33:08 -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 3176d4e358..585684567c 100644 --- a/unsloth/kernels/rms_layernorm.py +++ b/unsloth/kernels/rms_layernorm.py @@ -206,9 +206,9 @@ class Fast_RMS_Layernorm(torch.autograd.Function): pass -def fast_rms_layernorm(layernorm, X, gemma = False): +def fast_rms_layernorm(layernorm, X, gemma : bool = False): W = layernorm.weight - eps = layernorm.variance_epsilon if \ + eps : float = layernorm.variance_epsilon if \ hasattr(layernorm, "variance_epsilon") \ else layernorm.eps out = Fast_RMS_Layernorm.apply(X, W, eps, gemma)