mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
e9cd781681
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/15030 Reviewed By: bddppq Differential Revision: D13408998 fbshipit-source-id: 9eb675e09fbc4829eab34df7aa660a0590816feb
39 lines
833 B
C++
39 lines
833 B
C++
#ifndef CAFFE2_NET_ASYNC_TASK_H
|
|
#define CAFFE2_NET_ASYNC_TASK_H
|
|
|
|
#include "caffe2/core/net_async_base.h"
|
|
#include "caffe2/core/net_async_task_future.h"
|
|
#include "caffe2/core/operator.h"
|
|
|
|
#include <vector>
|
|
|
|
namespace caffe2 {
|
|
|
|
// AsyncTask represents an asynchronous execution of a chain of ops.
|
|
class AsyncTask {
|
|
public:
|
|
AsyncTask(const std::vector<OperatorBase*>& ops);
|
|
|
|
bool Run(const ExecutionOptions& options);
|
|
|
|
void Reset();
|
|
|
|
DeviceOption GetDeviceOption() const;
|
|
|
|
AsyncTaskFuture& GetFuture();
|
|
const AsyncTaskFuture& GetFuture() const;
|
|
|
|
private:
|
|
void handleChainError(
|
|
OperatorBase* op,
|
|
const char* err_msg,
|
|
bool save_exception = false);
|
|
|
|
std::vector<OperatorBase*> ops_;
|
|
DeviceOption device_option_;
|
|
AsyncTaskFuture future_;
|
|
};
|
|
|
|
} // namespace caffe2
|
|
|
|
#endif // CAFFE2_NET_ASYNC_TASK_H
|