srgan: enable model to run with gpu

models/SRGAN/srgan.py

@@ -75,12 +75,13 @@ class GeneratorResnet(nn.Module):
     
 if __name__ == '__main__':
     current_dir = os.path.dirname(os.path.realpath(__file__))
-
+    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     model = GeneratorResnet()
-    model = torch.load(current_dir + '/srgan_checkpoint.pth', map_location=torch.device('cpu'))
+    model = torch.load(current_dir + '/srgan_checkpoint.pth', map_location=torch.device('cpu')).to(DEVICE)
     model.eval()
     with torch.no_grad():
         input_image = Image.open('images/demo.png')
         input_image = ToTensor()(input_image).unsqueeze(0)
+        input_image = input_image.to(DEVICE)
         output_image = model.test(input_image)
         print(output_image.max())