Update g_retriever.py

g_retriever.py

@@ -42,7 +42,7 @@ class GRetrieverModel(LlamaForCausalLM):
         # mean pooling
         g_embeds = global_mean_pool(n_embeds, graph.batch.to(n_embeds.device))
 
-        return g_embeds.to(self.model.device)
+        return g_embeds
     
     @wraps(LlamaForCausalLM.forward)
     def forward(
@@ -76,7 +76,7 @@ class GRetrieverModel(LlamaForCausalLM):
             
             # encode graph
             graph_embeds = self.encode_graphs(graph)
-            graph_embeds = self.projector(graph_embeds)
+            graph_embeds = self.projector(graph_embeds).to(self.model.device)
             
             # prepare for reserved ids (bos+graph)
             non_tokenized_ids = (inputs == -1).nonzero()