enh: torch mps support

This commit is contained in:
Timothy Jaeryang Baek 2025-01-05 00:27:44 -08:00
parent 614379a427
commit d477f73c7e

View File

@ -54,6 +54,9 @@ else:
DEVICE_TYPE = "cpu"
if torch.backends.mps.is_available() and torch.backends.mps.is_built():
DEVICE_TYPE = "mps"
####################################
# LOGGING
####################################