|
@@ -58,9 +58,9 @@ for files in os.listdir(opt.input_dir):
|
|
|
# preprocess, (-1, 1)
|
|
|
input_image = -1 + 2 * input_image
|
|
|
if opt.gpu > -1:
|
|
|
- input_image = Variable(input_image).cuda()
|
|
|
+ input_image = Variable(input_image, volatile=True).cuda()
|
|
|
else:
|
|
|
- input_image = Variable(input_image).float()
|
|
|
+ input_image = Variable(input_image, volatile=True).float()
|
|
|
# forward
|
|
|
output_image = model(input_image)
|
|
|
output_image = output_image[0]
|