Update main.py

main.py

@@ -26,7 +26,7 @@ class Model(object):
     def load_model(self, weight_path=WEIGHT_PATH):
         if not os.path.isfile("./weights/best_weight.pth"):
             response = wget.download("https://raw.githubusercontent.com/hungnguyen2611/super-resolution/master/weights/best_weight.pth", "./weights/best_weight.pth")
-        weight = torch.load(weight_path)
+        weight = torch.load(weight_path, map_location=torch.device(DEVICE))
         print("[LOADING] Loading encoder...")
         self.model_Enc.load_state_dict(weight['model_Enc'])
         print("[LOADING] Loading decoder...")
@@ -74,8 +74,7 @@ if __name__ == '__main__':
     title = "Super-Resolution Demo USR-DA Unofficial 🚀🚀🔥"
     description = ''' 
 <br>
-**This Demo expects low-quality and low-resolution images**
-**We are looking for collaborators! Collaborator** 
+**This Demo expects low-quality and low-resolution images, better visual on real-world images**
 </br>
 '''
     article = "<p style='text-align: center'><a href='https://openaccess.thecvf.com/content/ICCV2021/papers/Wang_Unsupervised_Real-World_Super-Resolution_A_Domain_Adaptation_Perspective_ICCV_2021_paper.pdf' target='_blank'>Unsupervised Real-World Super-Resolution: A Domain Adaptation Perspective</a> | <a href='https://github.com/hungnguyen2611/super-resolution.git' target='_blank'>Github Repo</a></p>"