From 333d3d9a51c013860193703c581ced889a5ec2aa Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Wed, 6 Mar 2024 04:52:59 +1100 Subject: [PATCH] Update gemma.py --- unsloth/models/gemma.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index ba7aac7d31..f4bb6d877e 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -156,8 +156,8 @@ def GemmaModel_fast_forward_inference( hidden_states = self.embed_tokens(input_ids) # 3072**0.5 = 55.5000 in bfloat16, whilst 55.4256 in float32 # 2048**0.5 = 45.2500 in bfloat16, whilst 45.2548 in float32 - # inputs_embeds *= torch.tensor(math_sqrt(self.config.hidden_size), dtype = inputs_embeds.dtype) - inputs_embeds *= math_sqrt(self.config.hidden_size) + # hidden_states *= torch.tensor(math_sqrt(self.config.hidden_size), dtype = hidden_states.dtype) + hidden_states *= math_sqrt(self.config.hidden_size) next_decoder_cache = [] for idx, decoder_layer in enumerate(self.layers):