fix bug for compile related to dlpack

AlibabaPAI · Dec 11, 2024 · 069c89e · 069c89e
1 parent 4c87fa4
commit 069c89e
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/torch_xla/csrc/dl_convertor.cpp b/torch_xla/csrc/dl_convertor.cpp
@@ -138,9 +138,7 @@ DLManagedTensor* toDLPack(const at::Tensor& input) {
     auto external_ref = pjrt_buffer->AcquireExternalReference();
     XLA_CHECK_OK(external_ref.status());
     pack->external_reference = std::move(external_ref.value());
-    xla::PjRtFuture<> future = pjrt_buffer->GetReadyFuture();
-    absl::Status status = future.Await();
-    XLA_CHECK_OK(status);
+    XLA_CHECK_OK(pjrt_buffer->GetReadyFuture().Await());
   }
   pack->buffer_reference = pjrt_buffer;
 

diff --git a/torch_xla/csrc/unwrap_data.h b/torch_xla/csrc/unwrap_data.h
@@ -11,6 +11,9 @@
 
 namespace torch_xla {
 
+runtime::ComputationClient::DataPtr UnwrapXlaData(
+    const torch::lazy::BackendDataPtr& data);
+
 std::vector<runtime::ComputationClient::DataPtr> UnwrapXlaData(
     absl::Span<const torch::lazy::BackendDataPtr> datas);