biwen.lbw yingda.chen 3 years ago
parent
commit
dac2167b7a
1 changed files with 3 additions and 1 deletions
  1. +3
    -1
      modelscope/pipelines/cv/skin_retouching_pipeline.py

+ 3
- 1
modelscope/pipelines/cv/skin_retouching_pipeline.py View File

@@ -44,8 +44,10 @@ class SkinRetouchingPipeline(Pipeline):
"""
super().__init__(model=model)

if device == 'gpu':
if torch.cuda.is_available() and device == 'gpu':
device = 'cuda'
else:
device = 'cpu'
model_path = os.path.join(self.model, ModelFile.TORCH_MODEL_FILE)
detector_model_path = os.path.join(
self.model, 'retinaface_resnet50_2020-07-20_old_torch.pth')


Loading…
Cancel
Save