diff --git a/DeepPurpose/DTI.py b/DeepPurpose/DTI.py
index de625f7..2139c4f 100644
--- 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
index 7e9d946..06af574 100644
--- 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
index 33ca28d..7b35d22 100644
--- 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="kexinhuang@hsph.harvard.edu",