Update demo_img.py

demo_img.py

@@ -9,7 +9,7 @@ from networks.amtl import Model as AMTL
 from networks.amtg import Model as AMTG
 from utils import img2tensor, tensor2img, InputPadder
 
-device = torch.device('cpu' if torch.cuda.is_available() else 'cpu')
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model_dict = {
     'AMT-S': AMTS, 'AMT-L': AMTL, 'AMT-G': AMTG
 }