Update lance_ai_model.py

lance_ai_model.py

@@ -36,7 +36,7 @@ class LanceAI(PreTrainedModel, GenerationMixin):
 
         self.init_weights()
 
-    def forward(self, input_ids=None, attention_mask=None, labels=None, inputs_embeds=None, return_dict=True):
+    def forward(self, input_ids=None, attention_mask=None, labels=None, inputs_embeds=None, return_dict=True, use_cache=False):
         embeddings = self.embedding(input_ids) if inputs_embeds is None else inputs_embeds
         encoder_output = self.encoder(embeddings)
         decoder_output = self.decoder(embeddings, encoder_output)