0.1.3 fix a bug to integrate gpus for dgl models

kexinhuang12345 · May 14, 2021 · b77535a · b77535a
1 parent fc750f7
commit b77535a
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/DeepPurpose/encoders.py b/DeepPurpose/encoders.py
@@ -351,6 +351,7 @@ def __init__(self, in_feats, hidden_feats=None, max_degree = None, activation=No
 		self.transform = nn.Linear(predictor_hidden_size * 2, predictor_dim)
 
 	def forward(self, bg):
+		bg = bg.to(device)        
 		feats = bg.ndata.pop('h') 
 		node_feats = self.gnn(bg, feats)
 		node_feats = self.node_to_graph(node_feats)
@@ -373,6 +374,7 @@ def __init__(self, predictor_dim=None):
 		self.transform = nn.Linear(300, predictor_dim)
 
 	def forward(self, bg):
+		bg = bg.to(device)                
 		node_feats = [
 			bg.ndata.pop('atomic_number'),
 			bg.ndata.pop('chirality_type')
@@ -400,6 +402,7 @@ def __init__(self, predictor_dim=None):
 		self.transform = nn.Linear(300, predictor_dim)
 
 	def forward(self, bg):
+		bg = bg.to(device)                
 		node_feats = [
 			bg.ndata.pop('atomic_number'),
 			bg.ndata.pop('chirality_type')
@@ -432,6 +435,7 @@ def __init__(self, node_feat_size, edge_feat_size, num_layers = 2, num_timesteps
 		self.transform = nn.Linear(graph_feat_size, predictor_dim)
 
 	def forward(self, bg):
+		bg = bg.to(device)                
 		node_feats = bg.ndata.pop('h')
 		edge_feats = bg.edata.pop('e')
 

diff --git a/setup.py b/setup.py
@@ -15,7 +15,7 @@ def readme():
     name="DeepPurpose", 
     packages = ['DeepPurpose'],
     package_data={'DeepPurpose': ['ESPF/*']},
-    version="0.1.2",
+    version="0.1.3",
     author="Kexin Huang, Tianfan Fu",
     license="BSD-3-Clause",
     author_email="[email protected]",