Fix loading of scalar tensors. (#1234)

* Fix loading of scalar tensors. * Add regression test
mlverse · Jan 16, 2025 · b31a5fc · b31a5fc
1 parent eb05688
commit b31a5fc
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/R/tensor.R b/R/tensor.R
@@ -508,6 +508,9 @@ tensor_to_complex <- function(x) {
 #'
 #' @export
 torch_tensor_from_buffer <- function(buffer, shape, dtype = "float") {
+  if (!is.integer(shape)) {
+    shape <- as.integer(shape)
+  }
   cpp_tensor_from_buffer(buffer, shape, list(dtype=dtype))
 }
 

diff --git a/tests/testthat/test-save.R b/tests/testthat/test-save.R
@@ -424,4 +424,11 @@ test_that("can save a complex tensor", {
 
   expect_true(torch_allclose(x$real, z$real))
   expect_true(torch_allclose(x$imag, z$imag))
+})
+
+test_that("can load a scalar tensor", {
+  x <- torch_scalar_tensor(1)
+  k <- torch_serialize(x)
+  y <- torch_load(k)
+  expect_true(torch_allclose(x, y))
 })