Merge pull request ORNL#71 from allaffa/remap_all_fields_of_data_to_gpus

data.batch is remapped to the same device as data.x
RylieWeaver · Dec 16, 2021 · 9e8cae6 · 9e8cae6
2 parents 56a581b + cae3816
commit 9e8cae6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/hydragnn/models/Base.py b/hydragnn/models/Base.py
@@ -220,7 +220,7 @@ def forward(self, data):
 
         #### multi-head decoder part####
         # shared dense layers for graph level output
-        x_graph = global_mean_pool(x, batch)
+        x_graph = global_mean_pool(x, batch.to(x.device))
         outputs = []
         for head_dim, headloc, type_head in zip(
             self.head_dims, self.heads_NN, self.head_type