diff --git a/src/pie_modules/models/base_models/bart_as_pointer_network.py b/src/pie_modules/models/base_models/bart_as_pointer_network.py
index ffb1ce9bf..9d9a3cc60 100644
--- a/src/pie_modules/models/base_models/bart_as_pointer_network.py
+++ b/src/pie_modules/models/base_models/bart_as_pointer_network.py
@@ -7,7 +7,7 @@
 from transformers import BartConfig, BartModel, BartPreTrainedModel
 from transformers.modeling_outputs import Seq2SeqLMOutput, Seq2SeqModelOutput
 from transformers.models.bart.modeling_bart import shift_tokens_right
-from transformers.utils import logging
+from transformers.utils import ModelOutput, logging
 
 from pie_modules.models.base_models.bart_with_decoder_position_ids import (
     BartModelWithDecoderPositionIds,
@@ -58,6 +58,7 @@ def __init__(
         shared_decay: Optional[float] = None,
         encoder_layer_norm_decay: Optional[float] = 0.001,
         decoder_layer_norm_decay: Optional[float] = None,
+        use_key_value_cache: bool = True,
         # other BartConfig parameters
         **kwargs,
     ):
@@ -82,6 +83,8 @@ def __init__(
         self.encoder_layer_norm_decay = encoder_layer_norm_decay
         self.decoder_layer_norm_decay = decoder_layer_norm_decay
 
+        self.use_key_value_cache = use_key_value_cache
+
 
 class BartAsPointerNetwork(BartPreTrainedModel):
     config_class = BartAsPointerNetworkConfig
@@ -311,6 +314,14 @@ def forward(
             encoder_attentions=outputs.encoder_attentions,
         )
 
+    def _extract_past_from_model_output(
+        self, outputs: ModelOutput, standardize_cache_format: bool = False
+    ):
+        if not self.config.use_key_value_cache:
+            return None
+        else:
+            return super()._extract_past_from_model_output(outputs, standardize_cache_format)
+
     def prepare_inputs_for_generation(
         self,
         decoder_input_ids,