From f61612206cee9e00011ba9945be771ddc1673092 Mon Sep 17 00:00:00 2001 From: Junjie Bai Date: Thu, 13 Dec 2018 11:46:03 -0800 Subject: [PATCH] Fix old tensor OutputTensorCopyFrom usage in ImageInput operator (#15094) Summary: cc jerryzh168 Pull Request resolved: https://github.com/pytorch/pytorch/pull/15094 Differential Revision: D13451898 Pulled By: bddppq fbshipit-source-id: 27906be62fb88aaa13c257441a2e35a285b445ee --- caffe2/image/image_input_op.h | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/caffe2/image/image_input_op.h b/caffe2/image/image_input_op.h index 321d017..f5db8b3 100644 --- a/caffe2/image/image_input_op.h +++ b/caffe2/image/image_input_op.h @@ -1243,12 +1243,14 @@ bool ImageInputOp::CopyPrefetched() { // Note(jiayq): The if statement below should be optimized away by the // compiler since std::is_same is a constexpr. if (std::is_same::value) { - OperatorBase::OutputTensorCopyFrom(0, options, prefetched_image_, &context_); - OperatorBase::OutputTensorCopyFrom(1, options, prefetched_label_, &context_); + OperatorBase::OutputTensorCopyFrom( + 0, options, prefetched_image_, /* async */ true); + OperatorBase::OutputTensorCopyFrom( + 1, options, prefetched_label_, /* async */ true); for (int i = 2; i < OutputSize(); ++i) { OperatorBase::OutputTensorCopyFrom( - i, options, prefetched_additional_outputs_[i - 2], &context_); + i, options, prefetched_additional_outputs_[i - 2], /* async */ true); } } else { // TODO: support color jitter and color lighting in gpu_transform @@ -1276,14 +1278,17 @@ bool ImageInputOp::CopyPrefetched() { } else { OperatorBase::OutputTensorCopyFrom( - 0, type, prefetched_image_on_device_, &context_); + 0, type, prefetched_image_on_device_, /* async */ true); } OperatorBase::OutputTensorCopyFrom( - 1, type, prefetched_label_on_device_, &context_); + 1, type, prefetched_label_on_device_, /* async */ true); for (int i = 2; i < OutputSize(); ++i) { OperatorBase::OutputTensorCopyFrom( - i, type, prefetched_additional_outputs_on_device_[i - 2], &context_); + i, + type, + prefetched_additional_outputs_on_device_[i - 2], + /* async */ true); } } return true; -- 2.7.4