Эх сурвалжийг харах

Combines PR-#26 and #28 (#29)

* support CPU for inference_images.py

* introduce --num-works option
Shi Jin 4 жил өмнө
parent
commit
3d819586bf
1 өөрчлөгдсөн 1 нэмэгдсэн , 1 устгасан
  1. 1 1
      inference_images.py

+ 1 - 1
inference_images.py

@@ -89,7 +89,7 @@ if args.model_type == 'mattingrefine':
         args.model_refine_kernel_size)
 
 model = model.to(device).eval()
-model.load_state_dict(torch.load(args.model_checkpoint), strict=False)
+model.load_state_dict(torch.load(args.model_checkpoint, map_location=device), strict=False)
 
 
 # Load images