pytorch/caffe2/operators/free_op.h
Sebastian Messmer b0d3165cc8 refactor caffe2 operator constructors - 3/9 (#17084)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/17084

clangr codemod

Reviewed By: ezyang

Differential Revision: D14078507

fbshipit-source-id: ed02d772890b30196302b6830f541f054b7e95c8
2019-02-28 14:13:17 -08:00

30 lines
777 B
C++

#ifndef CAFFE2_OPERATORS_FREE_OP_H_
#define CAFFE2_OPERATORS_FREE_OP_H_
#include "caffe2/core/context.h"
#include "caffe2/core/operator.h"
namespace caffe2 {
// FreeOp frees the content of the output blob. We allow it to take in input
// blobs purely for the reason that it can "wait" on the input blobs to be
// produced by some of the earlier operators before a free is called.
template <class Context>
class FreeOp : public Operator<Context> {
public:
template <class... Args>
explicit FreeOp(Args&&... args)
: Operator<Context>(std::forward<Args>(args)...) {}
bool RunOnDevice() override {
for (Blob* output : OperatorBase::Outputs()) {
output->Reset();
}
return true;
}
};
} // namespace caffe2
#endif // CAFFE2_OPERATORS_FREE_OP_H_