Update g_retriever.py

g_retriever.py

@@ -80,10 +80,12 @@ class GRetrieverModel(LlamaForCausalLM):
             
             # embed inputs
             inputs[non_tokenized_shape] = self.config.pad_token_id
-            inputs_embeds = self.get_input_embeddings()(inputs)
+            temp_inputs_embeds = self.get_input_embeddings()(inputs)
             non_tokenized_embeds = torch.cat([bos_embeds.repeat(len(inputs), 1, 1), graph_embeds.unsqueeze(1)], dim=1)
             
             # replace reserved ids with bos+graph
+            temp_inputs_embeds[non_tokenized_shape] = non_tokenized_embeds.view(len(non_tokenized_ids), -1)
+            inputs_embeds = temp_inputs_embeds.clone()
             inputs_embeds[non_tokenized_shape] = non_tokenized_embeds.view(len(non_tokenized_ids), -1)
         
         else: