Added Torch::Backends::MPS.available? method

ankane · Apr 13, 2023 · b0174ec · b0174ec
1 parent d8230d7
commit b0174ec
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 2 deletions.
diff --git a/examples/mnist/main.rb b/examples/mnist/main.rb
@@ -81,8 +81,15 @@ def test(model, device, test_loader)
 
 Torch.manual_seed(seed)
 
-use_cuda = Torch::CUDA.available?
-device = Torch.device(use_cuda ? "cuda" : "cpu")
+device_type =
+  if Torch::CUDA.available?
+    "cuda"
+  elsif Torch::Backends::MPS.available?
+    "mps"
+  else
+    "cpu"
+  end
+device = Torch.device(device_type)
 puts "Device type: #{device.type}"
 
 root = File.join(__dir__, "data")

diff --git a/ext/torch/backends.cpp b/ext/torch/backends.cpp
@@ -14,4 +14,8 @@ void init_backends(Rice::Module& m) {
   Rice::define_module_under(rb_mBackends, "MKL")
     .add_handler<torch::Error>(handle_error)
     .define_singleton_function("available?", &torch::hasMKL);
+
+  Rice::define_module_under(rb_mBackends, "MPS")
+    .add_handler<torch::Error>(handle_error)
+    .define_singleton_function("available?", &torch::hasMPS);
 }
diff --git a/test/backends_test.rb b/test/backends_test.rb
@@ -8,4 +8,8 @@ def test_openmp
   def test_mkl
     Torch::Backends::MKL.available?
   end
+
+  def test_mps
+    Torch::Backends::MPS.available?
+  end
 end