Fix XLA tensor storage device by using XlaDeviceToAtenDevice. (#5743)

* Add test. * Convert storage device using `XlaDeviceToAtenDevice`.
pytorch · Jan 12, 2024 · 0d83c1b · 0d83c1b
1 parent 1b54693
commit 0d83c1b
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/test/cpp/test_aten_xla_tensor_1.cpp b/test/cpp/test_aten_xla_tensor_1.cpp
@@ -23,6 +23,16 @@ class AtenXlaTensorTest : public AtenXlaTensorTestBase {};
 
 }  // namespace
 
+TEST_F(AtenXlaTensorTest, TestStorage) {
+  torch::Tensor a = torch::tensor({0.0});
+  ForEachDevice([&](const torch::Device& device) {
+    torch::Tensor xla_a = CopyToDevice(a, device);
+    XLATensorPtr xla_tensor_a = bridge::GetXlaTensor(xla_a);
+    EXPECT_EQ(xla_a.device(), xla_tensor_a->Storage().device());
+    AllClose(a, xla_a);
+  });
+}
+
 TEST_F(AtenXlaTensorTest, TestEmpty) {
   torch::Tensor a = torch::zeros({2, 2}, torch::TensorOptions(torch::kFloat));
   ForEachDevice([&](const torch::Device& device) {

diff --git a/torch_xla/csrc/tensor.cpp b/torch_xla/csrc/tensor.cpp
@@ -19,6 +19,7 @@
 #include <stdexcept>
 #include <unordered_set>
 
+#include "torch_xla/csrc/aten_xla_bridge.h"
 #include "torch_xla/csrc/debug_util.h"
 #include "torch_xla/csrc/helpers.h"
 #include "torch_xla/csrc/layout_manager.h"
@@ -139,9 +140,9 @@ XLATensor::XLATensor(std::shared_ptr<View> view,
 XLATensor::XLATensor(std::shared_ptr<Data> data)
     : torch::lazy::LazyTensor(data),
       data_(std::move(data)),
-      storage_(c10::Storage(
-          {}, 0,
-          c10::DataPtr(nullptr, backendDeviceToAtenDevice(data_->device)))) {}
+      storage_(c10::Storage({}, 0,
+                            c10::DataPtr(nullptr, bridge::XlaDeviceToAtenDevice(
+                                                      data_->device)))) {}
 
 auto XLATensor::data() const -> const std::shared_ptr<Data>& {
   XLA_CHECK(data_ != nullptr) << "Trying to access a null cursor";