Modifications based on PR #91 review

Author: mickaelseznec

main.py

@@ -163,10 +163,7 @@ def main():
         network_data = None
         print("=> creating model '{}'".format(args.arch))
 
-    if device.type == "cuda":
-        model = models.__dict__[args.arch](network_data).cuda()
-    else:
-        model = models.__dict__[args.arch](network_data).cpu()
+    model = models.__dict__[args.arch](network_data).to(device)
 
     assert(args.solver in ['adam', 'sgd'])
     print('=> setting {} solver'.format(args.solver))

run_inference.py

@@ -71,7 +71,7 @@ def main():
     for ext in args.img_exts:
         test_files = data_dir.files('*1.{}'.format(ext))
         for file in test_files:
-            img_pair = file.parent / (file.basename().splitext()[0][:-1] + '2.{}'.format(ext))
+            img_pair = file.parent / (file.stem[:-1] + '2.{}'.format(ext))
             if img_pair.isfile():
                 img_pairs.append([file, img_pair])
 
@@ -103,7 +103,7 @@ def main():
         if args.upsampling is not None:
             output = F.interpolate(output, size=img1.size()[-2:], mode=args.upsampling, align_corners=False)
         for suffix, flow_output in zip(['flow', 'inv_flow'], output):
-            filename = save_path/'{}{}'.format(img1_file.basename().splitext()[0][:-1], suffix)
+            filename = save_path/'{}{}'.format(img1_file.stem[:-1], suffix)
             if args.output_value in['vis', 'both']:
                 rgb_flow = flow2rgb(args.div_flow * flow_output, max_value=args.max_flow)
                 to_save = (rgb_flow * 255).astype(np.uint8).transpose(1,2,0)