fix comment

pytorch · May 1, 2024 · 3d7e463 · 3d7e463
1 parent 8dcf30f
commit 3d7e463
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/torch_xla/csrc/init_python_bindings.cpp b/torch_xla/csrc/init_python_bindings.cpp
@@ -2482,10 +2482,12 @@ void InitXlaModuleBindings(py::module m) {
               return runtime::GetComputationClient()->UnsafeBufferPointer(
                   UnwrapXlaData(data));
             } else {
-              XLA_ERROR() << "Could not get the buffer pointer.";
+              XLA_ERROR() << "Could not get the buffer pointer for XLATensor "
+                             "with IR that's not DeviceData";
             }
           }
-          XLA_ERROR() << "Could not get the buffer pointer.";
+          XLA_ERROR() << "Could not get the buffer pointer for XLATensor "
+                         "without a data handle or an IR.";
         });
 
   // -------------Dynamo Integration API Start-------------------------

diff --git a/torch_xla/csrc/runtime/pjrt_computation_client.cc b/torch_xla/csrc/runtime/pjrt_computation_client.cc
@@ -460,8 +460,9 @@ std::vector<ComputationClient::DataPtr> PjRtComputationClient::ReshardData(
 
 std::uintptr_t PjRtComputationClient::UnsafeBufferPointer(
     const DataPtr handle) {
-  std::shared_ptr<PjRtData> pjrt_data = ReplicateShardedData(handle);
-  XLA_CHECK(pjrt_data);
+  std::shared_ptr<PjRtData> pjrt_data =
+      std::dynamic_pointer_cast<PjRtData>(handle) XLA_CHECK(pjrt_data)
+      << "handle must be PjRtData, got " << handle->ToString();
   xla::StatusOr<std::uintptr_t> ptr =
       client_->UnsafeBufferPointer(pjrt_data->buffer.get());
   XLA_CHECK(ptr.ok());