Skip to content

Commit

Permalink
GPU support on Apple Silicon (#30)
Browse files Browse the repository at this point in the history
* 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>
  • Loading branch information
3 people committed May 9, 2023
1 parent b0cac91 commit ba1b0d9
Showing 1 changed file with 3 additions and 0 deletions.
3 changes: 3 additions & 0 deletions manga_ocr/ocr.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,9 @@ def __init__(self, pretrained_model_name_or_path='kha-white/manga-ocr-base', for
if not force_cpu and torch.cuda.is_available():
logger.info('Using CUDA')
self.model.cuda()
if not force_cpu and torch.backends.mps.is_available():
logger.info('Using MPS')
self.model.to('mps')
else:
logger.info('Using CPU')

Expand Down

0 comments on commit ba1b0d9

Please sign in to comment.