thenAsync Class — pytorch Architecture

Source Code

aten/src/ATen/core/ivalue_inl.h lines 1099–1126

  template <typename T>
  c10::intrusive_ptr<Future> thenAsync(T callback, TypePtr type) {
    static_assert(
        std::is_invocable_r_v<c10::intrusive_ptr<Future>, T, Future&>,
        "The callback must have signature c10::intrusive_ptr<Future>(Future&)");

    auto childFut = createInstance(std::move(type));
    addCallback(
        [childFut, cb = std::move(callback)](Future& parentFut) mutable {
          c10::intrusive_ptr<Future> intermediateFut;
          try {
            intermediateFut = cb(parentFut);
          } catch (std::exception&) {
            childFut->setError(std::current_exception());
            return;
          }
          intermediateFut->addCallback(
              [childFut = std::move(childFut)](Future& intermediateFut) {
                if (intermediateFut.hasError()) {
                  childFut->setError(intermediateFut.exception_ptr());
                } else {
                  childFut->markCompleted(
                      intermediateFut.value(), intermediateFut.storages());
                }
              });
        });
    return childFut;
  }