Преглед изворни кода

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

* support CPU for inference_images.py

* introduce --num-works option
Shi Jin пре 4 година
родитељ
комит
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