Yijun Li 6 lat temu
rodzic
commit
760b975e55
1 zmienionych plików z 2 dodań i 2 usunięć
  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]