try that

mlverse · Feb 28, 2024 · 8d4b1c3 · 8d4b1c3
1 parent 03d04e1
commit 8d4b1c3
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/R/accelerator.R b/R/accelerator.R
@@ -93,7 +93,7 @@ LuzAcceleratorState <- R6::R6Class(
 
       if (torch::cuda_is_available())
         paste0("cuda:", index)
-      else if (torch::backends_mps_is_available())
+      else if (can_use_mps())
         "mps"
       else
         "cpu"

diff --git a/R/module.R b/R/module.R
@@ -603,8 +603,13 @@ get_metrics.luz_module_evaluation <- function(object, ...) {
   res[, c("metric", "value")]
 }
 
+can_use_mps <- function() {
+  arch <- Sys.info()["machine"]
+  "arm64" %in% arch && torch::backends_mps_is_available()
+}
+
 enable_mps_fallback <- function() {
-  if (!torch::backends_mps_is_available())
+  if (!can_use_mps())
     return(invisible(NULL))
 
   fallback <- Sys.getenv("PYTORCH_ENABLE_MPS_FALLBACK", unset = "")