then Class — pytorch Architecture

Source Code

aten/src/ATen/core/ivalue_inl.h lines 1072–1097

  template <typename T>
  c10::intrusive_ptr<Future> then(T callback, TypePtr type) {
    using IValueWithStorages = std::tuple<IValue, std::vector<WeakStorage>>;
    static_assert(
        std::disjunction_v<
            std::is_invocable_r<IValue, T, Future&>,
            std::is_invocable_r<IValueWithStorages, T, Future&>>,
        "The callback must have signature IValue(Future&) or "
        "std::tuple<IValue, std::vector<Storage>>(Future&)");

    auto childFut = createInstance(::std::move(type));
    addCallback([childFut,
                 cb = std::move(callback)](Future& parentFut) {
      try {
        if constexpr (::std::is_convertible_v<typename std::invoke_result_t<T &&, Future&>, IValueWithStorages>) {
          auto [ivalue, storages] = cb(parentFut);
          childFut->markCompleted(::std::move(ivalue), ::std::move(storages));
        } else {
          childFut->markCompleted(cb(parentFut));
        }
      } catch (std::exception&) {
        childFut->setError(std::current_exception());
      }
    });
    return childFut;
  }