Merge pull request #7 from SCAuFish/main

Support Runtime on Apple Silicon
apple · Oct 28, 2024 · 32d5b01 · 32d5b01
2 parents f9e87fc + ada9e77
commit 32d5b01
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/ml_mdm/clis/generate_sample.py b/ml_mdm/clis/generate_sample.py
@@ -20,7 +20,13 @@
 from ml_mdm.config import get_arguments, get_model, get_pipeline
 from ml_mdm.language_models import factory
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device(
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps"
+    if torch.backends.mps.is_available()
+    else "cpu"
+)
 
 # Note that it is called add_arguments, not add_argument.
 logging.basicConfig(

diff --git a/pyproject.toml b/pyproject.toml
@@ -23,7 +23,7 @@ dependencies = [
     "imageio[ffmpeg]",
     "matplotlib",
     "mlx-data",
-    "numpy",
+    "numpy<2",
     "pytorch-model-summary",
     "rotary-embedding-torch",
     "simple-parsing==0.1.5",