attention mask

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

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,