#ifndef CAFFE2_OPERATORS_COPY_OP_H_ #define CAFFE2_OPERATORS_COPY_OP_H_ #include "caffe2/core/context.h" #include "caffe2/core/operator.h" namespace caffe2 { template class CopyOp : public Operator { public: USE_OPERATOR_CONTEXT_FUNCTIONS; USE_SIMPLE_CTOR_DTOR(CopyOp) bool RunOnDevice() override { auto& input = this->template Input(0, SrcContext::GetDeviceType()); auto* output = this->template Output(0, DstContext::GetDeviceType()); output->ResizeLike(input); this->context_.template CopyItems( input.dtype(), input.numel(), input.raw_data(), output->raw_mutable_data(input.dtype())); return true; } }; template class CopyOnDeviceLikeOp : public CopyOp { public: template explicit CopyOnDeviceLikeOp(Args&&... args) : CopyOp(std::forward(args)...) {} }; } // namespace caffe2 #endif // CAFFE2_OPERATORS_COPY_OP_H_