8000 Use `addressable_devices_` instead of `devices_` in case of the multi… · linux-on-ibm-z/tensorflow@7587767 · GitHub
[go: up one dir, main page]
Skip to content

Commit 7587767

Browse files
tensorflower-gardenertensorflower-gardener
committed
Use addressable_devices_ instead of devices_ in case of the multi-host environment.
PiperOrigin-RevId: 730551286
1 parent d615e26 commit 7587767

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

‎third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1238,7 +1238,7 @@ absl::Status PjRtStreamExecutorClient::DmaMap(void* data, size_t buffer_size) {
12381238
tsl::profiler::TraceMe trace_me("PjRtStreamExecutorClient::DmaMap");
12391239
TF_ASSIGN_OR_RETURN(
12401240
LocalDeviceState * local_device,
1241-
tensorflow::down_cast<PjRtStreamExecutorDevice*>(devices_[0])
1241+
tensorflow::down_cast<PjRtStreamExecutorDevice*>(addressable_devices_[0])
12421242
->GetLocalDeviceState());
12431243
bool success = local_device->compute_stream()->parent()->HostMemoryRegister(
12441244
data, buffer_size);
@@ -1255,7 +1255,7 @@ absl::Status PjRtStreamExecutorClient::DmaUnmap(void* data) {
12551255
tsl::profiler::TraceMe trace_me("PjRtStreamExecutorClient::DmaUnmap");
12561256
TF_ASSIGN_OR_RETURN(
12571257
LocalDeviceState * local_device,
1258-
tensorflow::down_cast<PjRtStreamExecutorDevice*>(devices_[0])
1258+
tensorflow::down_cast<PjRtStreamExecutorDevice*>(addressable_devices_[0])
12591259
->GetLocalDeviceState());
12601260
bool success =
12611261
local_device->compute_stream()->parent()->HostMemoryUnregister(data);

0 commit comments

Comments
 (0)
0