zack3d/pylingual
1
0
Fork 0
mirror of https://github.com/syssec-utd/pylingual.git synced 2026-05-10 18:39:03 -07:00
Code Issues Packages Projects Releases Wiki Activity

mps support

Browse Source
This commit is contained in:
Joel Flores
2025-06-26 17:37:08 -05:00
parent 3fae0294fb
commit 3ab3abccc1
1 changed files with 4 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pylingual/models.py
+4
View File
@@ -116,7 +116,11 @@ def load_models(config_file: Path = Path("pylingual/decompiler_config.yaml"), ve
mask_token="[MASK]",
)
if torch.cuda.is_available():
logger.info("Using CUDA GPU for models")
device = torch.device("cuda:0")
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
logger.info("Using MPS (Metal Performance Shaders) GPU for models")
device = torch.device("mps")
else:
logger.warning("Using CPU for models")
device = torch.device("cpu")