Update pjrt_computation_client.cc

pytorch · Dec 4, 2023 · 4a220e7 · 4a220e7
1 parent 339ad72
commit 4a220e7
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/torch_xla/csrc/runtime/pjrt_computation_client.cc b/torch_xla/csrc/runtime/pjrt_computation_client.cc
@@ -172,6 +172,7 @@ PjRtComputationClient::PjRtComputationClient() {
     options.num_nodes = global_world_size;
     options.allowed_devices = allowed_devices;
     options.platform_name = "gpu";
+    options.should_stage_host_to_device_transfers = true;
     options.kv_get = kv_get;
     options.kv_put = kv_put;
     client_ = std::move(xla::GetStreamExecutorGpuClient(options).value());