Merge branch 'main' into minor-upds

AllenInstitute · Oct 19, 2024 · 17e40ab · 17e40ab
2 parents 7b79c26 + 0942992
commit 17e40ab
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/deep_neurographs/machine_learning/heterograph_models.py b/src/deep_neurographs/machine_learning/heterograph_models.py
@@ -154,7 +154,7 @@ def forward(self, x_dict, edge_index_dict, edge_attr_dict):
 
         # Input - Edges
         for key, f in self.input_edges.items():
-            key = reformat_edge_key(key)
+            #key = reformat_edge_key(key)
             edge_attr_dict[key] = f(edge_attr_dict[key])
         edge_attr_dict = self.activation(edge_attr_dict)
 

diff --git a/src/deep_neurographs/utils/gnn_util.py b/src/deep_neurographs/utils/gnn_util.py
@@ -25,7 +25,7 @@ def get_inputs(data, device=None, is_multimodal=False):
     x = data.x_dict
     edge_index = data.edge_index_dict
     edge_attr = data.edge_attr_dict
-    if device and torch.cuda.is_available():
+    if "cuda" in device and torch.cuda.is_available():
         return toGPU(x), toGPU(edge_index), toGPU(edge_attr)
     else:
         return x, edge_index, edge_attr