oguzakif commited on
Commit
6c89e77
·
1 Parent(s): e7aba9f

torch device selector added

Browse files
Files changed (1) hide show
  1. FGT_codes/tool/video_inpainting.py +2 -1
FGT_codes/tool/video_inpainting.py CHANGED
@@ -415,7 +415,8 @@ def save_results(outdir, comp_frames):
415
 
416
 
417
  def video_inpainting(args, imgArr, imgMaskArr):
418
- device = torch.device("cuda:{}".format(args.gpu))
 
419
  print(args)
420
  if args.opt is not None:
421
  with open(args.opt, "r") as f:
 
415
 
416
 
417
  def video_inpainting(args, imgArr, imgMaskArr):
418
+ #device = torch.device("cuda:{}".format(args.gpu))
419
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
420
  print(args)
421
  if args.opt is not None:
422
  with open(args.opt, "r") as f: