Error when using --cuda in demo.py

[advaza]

Hi,
I'm trying to run inference using the demo.py code. I did minimal changes: enable to get any king of input image and save the output image instead of displaying it (link to code). I used really small image for evaluation. When I didn't add the cuda flag time is took for processing was 0.125365018845. When I added the --cuda flag, it got stuck without any error message.

OS: Linux 4.4.78
Python version: 2.7.12
PyTorch version: 0.4.0
CUDA version: 9.0.176

Thank you!
Adva