Преглед на файлове

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