Make sure we copy the tensor to cpu before copying to R. (#1080)

mlverse · Jul 11, 2023 · 79a0fe5 · 79a0fe5
1 parent 15b7803
commit 79a0fe5
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/R/tensor.R b/R/tensor.R
@@ -378,6 +378,8 @@ as.matrix.torch_tensor <- function(x, ...) {
 }
 
 as_array_impl <- function(x) {
+  # move tensor to cpu before copying to R
+  x <- x$cpu()
 
   if (x$is_complex()) {
     out <- complex(

diff --git a/tests/testthat/test-tensor.R b/tests/testthat/test-tensor.R
@@ -539,4 +539,13 @@ test_that("can create a buffer from a tensor", {
   y <- buffer_from_torch_tensor(x)
   z <- torch_tensor_from_buffer(y, shape = c(10, 10), dtype="float")
   expect_true(torch_allclose(x, z))
+})
+
+test_that("can copy a mps tensor", {
+  skip_if_not_m1_mac()
+  x <- array(runif(100), dim = c(10, 10))
+  y <- torch_tensor(x, device="mps")
+  x_ <- as.array(y)
+
+  expect_true(all.equal(x, x_, tolerance = 1e-5))
 })