From 71483e6c8d577f78fb0e130f9859e715a4786e50 Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Sat, 24 Feb 2024 17:17:26 +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 69c4535d44..d9172d2c3b 100644
--- a/unsloth/models/gemma.py
+++ b/unsloth/models/gemma.py
@@ -324,7 +324,7 @@ def GemmaModel_fast_forward(
             layer_outputs = torch.utils.checkpoint.checkpoint(
                 create_custom_forward(decoder_layer),
                 hidden_states,
-                # causal_mask,
+                causal_mask,
                 attention_mask,
                 position_ids,
                 use_reentrant=True,
@@ -333,7 +333,7 @@ def GemmaModel_fast_forward(
         else:
             layer_outputs = decoder_layer(
                 hidden_states,
-                # causal_mask=causal_mask,
+                causal_mask=causal_mask,
                 attention_mask=attention_mask,
                 position_ids=position_ids,
                 past_key_value=past_key_value,