diff --git a/examples/nlp/parameter_efficient_finetuning_of_gpt2_with_lora.py b/examples/nlp/parameter_efficient_finetuning_of_gpt2_with_lora.py
index 877f3adcae..64c9af58a9 100644
--- a/examples/nlp/parameter_efficient_finetuning_of_gpt2_with_lora.py
+++ b/examples/nlp/parameter_efficient_finetuning_of_gpt2_with_lora.py
@@ -587,6 +587,10 @@ def call(self, inputs):
     B_weights = value_lora_layer.B.kernel  # (1, 12, 64) (b, c, d)
     increment_weights = tf.einsum("ab,bcd->acd", A_weights, B_weights) * (ALPHA / RANK)
     value_lora_layer.original_layer.kernel.assign_add(increment_weights)
+    
+    # Put back in place the original layers with updated weights
+    self_attention_layer._query_dense = query_lora_layer.original_layer
+    self_attention_layer._value_dense = value_lora_layer.original_layer
 
 """
 We are now all set to generate text with our LoRA model :).