sync with xla

pytorch · Jul 18, 2024 · 21823e3 · 21823e3
1 parent eeded81
commit 21823e3
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/torch_xla/csrc/runtime/pjrt_compilation_client.cc b/torch_xla/csrc/runtime/pjrt_compilation_client.cc
@@ -91,7 +91,7 @@ PjRtCompilationClient::PjRtCompilationClient(
       env::kEnvTpuLibraryPath,
       sys_util::GetEnvString(env::kEnvInferredTpuLibraryPath, "libtpu.so"));
   XLA_CHECK_OK(pjrt::LoadPjrtPlugin("tpu", tpu_library_path).status());
-  xla::Status tpu_status = pjrt::InitializePjrtPlugin("tpu");
+  absl::Status tpu_status = pjrt::InitializePjrtPlugin("tpu");
   XLA_CHECK_OK(tpu_status);
 
   absl::flat_hash_map<std::string, xla::PjRtValueType> create_options = {};
@@ -626,7 +626,7 @@ int PjRtCompilationClient::GetNumProcesses() const {
 };
 
 const absl::flat_hash_map<
-    std::string, torch_xla::runtime::ComputationClient::DeviceAttribute>&
+    std::string, torch_xla::runtime::ComputationClient::DeviceAttribute>
 PjRtCompilationClient::GetDeviceAttributes(const std::string& device) {
   return PjRtCompilationClient::StringToPjRtDevice(device)->Attributes();
 }

diff --git a/torch_xla/csrc/runtime/pjrt_compilation_client.h b/torch_xla/csrc/runtime/pjrt_compilation_client.h
@@ -123,7 +123,7 @@ class PjRtCompilationClient : public ComputationClient {
   int GetNumProcesses() const override;
 
   const absl::flat_hash_map<
-      std::string, torch_xla::runtime::ComputationClient::DeviceAttribute>&
+      std::string, torch_xla::runtime::ComputationClient::DeviceAttribute>
   GetDeviceAttributes(const std::string& device) override;
 
   void SetReplicationDevices(