Merge branch 'main' into fix-all

apple · Nov 2, 2024 · 93b50a2 · 93b50a2
2 parents 9f9bd31 + 32d5b01
commit 93b50a2
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",