fix merge

pytorch · Nov 27, 2023 · f51ae47 · f51ae47
1 parent 80926b3
commit f51ae47
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/torch_xla/csrc/runtime/pjrt_computation_client.cc b/torch_xla/csrc/runtime/pjrt_computation_client.cc
@@ -641,14 +641,14 @@ PjRtComputationClient::ExecuteReplicated(
   const PjRtComputation& pjrt_computation =
       dynamic_cast<const PjRtComputation&>(computation);
 
+  std::vector<std::vector<xla::PjRtBuffer*>> argument_handles(
+      devices.size(), std::vector<xla::PjRtBuffer*>(arguments.size()));
   {
     tsl::profiler::TraceMe activity(
         "PjRtComputationClient::ExecuteReplicated_argument_handle",
         tsl::profiler::TraceMeLevel::kInfo);
 
     absl::BlockingCounter counter(arguments.size());
-    std::vector<std::vector<xla::PjRtBuffer*>> argument_handles(
-        devices.size(), std::vector<xla::PjRtBuffer*>(arguments.size()));
     // TODO: tune and document cost estimate
     pool_.ParallelFor(arguments.size(), 30000, [&](int64_t start, int64_t end) {
       tsl::profiler::TraceMe activity(
@@ -671,7 +671,6 @@ PjRtComputationClient::ExecuteReplicated(
         }
         counter.DecrementCount();
       };
-      thread::Schedule(std::move(buffer_converter));
     });
     counter.Wait();
   }