Update g_retriever.py

g_retriever.py

@@ -24,13 +24,13 @@ class GRetrieverModel(LlamaForCausalLM):
             num_layers=config.gnn_num_layers,
             dropout=config.gnn_dropout,
             num_heads=config.gnn_num_heads,
-        ).to(self.model.dtype).to(self.model.device)
+        ).to(self.model.dtype)
 
         self.projector = nn.Sequential(
             nn.Linear(config.gnn_hidden_dim, 2048),
             nn.Sigmoid(),
             nn.Linear(2048, self.get_input_embeddings().embedding_dim),
-        ).to(self.model.dtype).to(self.model.device)
+        ).to(self.model.dtype)
     
     def encode_graphs(self, graph):
         n_embeds, _ = self.graph_encoder(
@@ -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
+        return g_embeds.to(model.device)
     
     @wraps(LlamaForCausalLM.forward)
     def forward(