Update g_retriever.py

g_retriever.py

@@ -33,10 +33,14 @@ class GRetrieverModel(LlamaForCausalLM):
         ).to(self.model.dtype).to(self.model.device)
     
     def encode_graphs(self, graph):
-        n_embeds, _ = self.graph_encoder(graph.x.to(self.model.dtype), graph.edge_index.long(), graph.edge_attr.to(self.model.dtype))
+        n_embeds, _ = self.graph_encoder(
+            graph.x.to(self.model.dtype),
+            graph.edge_index.long(),
+            graph.edge_attr.to(self.model.dtype)
+        )
 
         # mean pooling
-        g_embeds = global_mean_pool(n_embeds, graph.batch)
+        g_embeds = global_mean_pool(n_embeds, graph.batch.to(n_embeds.device))
 
         return g_embeds