Merge pull request #46 from bkvie/patch-2

Update run_inference.py

Author: ClementPinard

run_inference.py

@@ -32,7 +32,7 @@
                     'which is 4 times downsampled. If set, will output full resolution flow map, with selected upsampling')
 parser.add_argument('--bidirectional', action='store_true', help='if set, will output invert flow (from 1 to 0) along with regular flow')
 
-
+@torch.no_grad()
 def main():
     global args, save_path
     args = parser.parse_args()
@@ -72,11 +72,11 @@ def main():
 
         img1 = input_transform(imread(img1_file))
         img2 = input_transform(imread(img2_file))
-        input_var = torch.autograd.Variable(torch.cat([img1, img2]).cuda(), volatile=True).unsqueeze(0)
+        input_var = torch.tensor(torch.cat([img1, img2]).cuda()).unsqueeze(0)
 
         if args.bidirectional:
             # feed inverted pair along with normal pair
-            inverted_input_var = torch.autograd.Variable(torch.cat([img2, img1],0).cuda(), volatile=True).unsqueeze(0)
+            inverted_input_var = torch.tensor(torch.cat([img2, img1],0).cuda()).unsqueeze(0)
             input_var = torch.cat([input_var, inverted_input_var])
 
         # compute output
@@ -85,7 +85,7 @@ def main():
             output = torch.nn.functional.upsample(output, size=img1.size()[-2:], mode=args.upsampling)
         for suffix, flow_output in zip(['flow', 'inv_flow'], output.data.cpu()):
             rgb_flow = flow2rgb(args.div_flow * flow_output.numpy(), max_value=args.max_flow)
-            to_save = (rgb_flow * 255).astype(np.uint8)
+            to_save = (rgb_flow * 255).astype(np.uint8).transpose(1,2,0)
             imsave(save_path/'{}{}.png'.format(img1_file.namebase[:-1], suffix), to_save)