TensorNest compute()

in src/cc/actorpool.cc [294:321]


  TensorNest compute(TensorNest tensors) {
    BatchPromise promise;
    auto future = promise.get_future();

    batching_queue_.enqueue({std::move(tensors), std::move(promise)});

    std::future_status status = future.wait_for(std::chrono::seconds(10 * 60));
    if (status != std::future_status::ready) {
      throw py::timeout_error("Compute timeout reached.");
    }

    const std::pair<std::shared_ptr<TensorNest>, int64_t> pair = [&] {
      try {
        return future.get();
      } catch (const std::future_error& e) {
        if (batching_queue_.is_closed() &&
            e.code() == std::future_errc::broken_promise) {
          throw ClosedBatchingQueue("Batching queue closed during compute");
        }
        throw;
      }
    }();

    return pair.first->map([batch_dim = batch_dim_,
                            batch_entry = pair.second](const torch::Tensor& t) {
      return t.slice(batch_dim, batch_entry, batch_entry + 1);
    });
  }