GPU support on Apple Silicon (#30)

* Added support for MPS on Apple Silicon

* Requires PyTorch 2.0.0 or greater

* Update requirements.txt

---------

Co-authored-by: Timo Mom <timo.mom@concedro.com>
Co-authored-by: Maciej Budyś <kha-white@mail.com>
This commit is contained in:
2023-05-09 20:10:15 +02:00
committed by GitHub
parent b0cac91777
commit ba1b0d94a8

View File

@@ -18,6 +18,9 @@ class MangaOcr:
if not force_cpu and torch.cuda.is_available(): if not force_cpu and torch.cuda.is_available():
logger.info('Using CUDA') logger.info('Using CUDA')
self.model.cuda() self.model.cuda()
if not force_cpu and torch.backends.mps.is_available():
logger.info('Using MPS')
self.model.to('mps')
else: else:
logger.info('Using CPU') logger.info('Using CPU')