Yijun Li 6 gadi atpakaļ
vecāks
revīzija
760b975e55
1 mainītis faili ar 2 papildinājumiem un 2 dzēšanām
  1. 2 2
      test.py

+ 2 - 2
test.py

@@ -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]