Conditional use of GPU on FastMtCnn

This commit is contained in:
aku-ato 2024-03-26 15:34:03 +01:00 committed by GitHub
parent 3d8733812b
commit 2b7c0a708e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -68,7 +68,9 @@ class FastMtCnnClient(Detector):
"Please install using 'pip install facenet-pytorch' " "Please install using 'pip install facenet-pytorch' "
) from e ) from e
face_detector = fast_mtcnn(device="cuda:0") device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
face_detector = fast_mtcnn(device=device)
return face_detector return face_detector