fix cuda device bug

kexinhuang12345 · Apr 23, 2021 · 3e61473 · 3e61473
1 parent fa3ee21
commit 3e61473
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 6 deletions.
diff --git a/DeepPurpose/DTI.py b/DeepPurpose/DTI.py
@@ -587,10 +587,7 @@ def load_pretrained(self, path):
 		if not os.path.exists(path):
 			os.makedirs(path)
 
-		if self.device[:4] == 'cuda':
-			state_dict = torch.load(path)
-		else:
-			state_dict = torch.load(path, map_location = torch.device('cpu'))
+		state_dict = torch.load(path, map_location = torch.device('cpu'))
 		# to support training from multi-gpus data-parallel:
 
 		if next(iter(state_dict))[:7] == 'module.':

diff --git a/README.md b/README.md
@@ -593,7 +593,7 @@ Checkout [Dataset Tutorial](DEMO/load_data_tutorial.ipynb).
 We provide more than 10 pretrained models. Please see [Pretraining Model Tutorial](DEMO/load_pretraining_models_tutorial.ipynb) on how to load them. It is as simple as 
 
 ```python
-from DeepPurpose import models
+from DeepPurpose import DTI as models
 net = models.model_pretrained(model = 'MPNN_CNN_DAVIS')
 or
 net = models.model_pretrained(FILE_PATH)

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.0",
+    version="0.1.1",
     author="Kexin Huang, Tianfan Fu",
     license="BSD-3-Clause",
     author_email="[email protected]",