diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index d5dd7833d5..69f75b32b9 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -486,6 +486,10 @@ def LlamaModel_fast_forward( if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) + inputs_embeds.requires_grad_(False) + inputs_embeds *= attention_mask.unsqueeze(0).transpose(0, 1).transpose(1, 2) + inputs_embeds.requires_grad_(True) + # Ignore attention_mask if attention_mask is None: padding_mask = None diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index fa61c1554f..2410572174 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -195,8 +195,6 @@ def MistralForCausalLM_fast_forward( return_dict = return_dict if return_dict is not None else self.config.use_return_dict # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - print(input_ids) - print(attention_mask) outputs = self.model( input_ids=input_ids, causal_mask=causal_mask,