Explorar o código

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

* support CPU for inference_images.py

* introduce --num-works option
Shi Jin %!s(int64=4) %!d(string=hai) anos
pai
achega
3d819586bf
Modificáronse 1 ficheiros con 1 adicións e 1 borrados
  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