Jelajahi Sumber

Update test.py

Yijun Li 6 tahun lalu
induk
melakukan
760b975e55
1 mengubah file dengan 2 tambahan dan 2 penghapusan
  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]