diff options
| author | Junjie Bai <bai@in.tum.de> | 2018-12-13 11:46:03 -0800 |
|---|---|---|
| committer | Facebook Github Bot <facebook-github-bot@users.noreply.github.com> | 2018-12-13 11:48:19 -0800 |
| commit | f61612206cee9e00011ba9945be771ddc1673092 (patch) | |
| tree | cb8e1adf63bc62d90a57eb765d10a85fc3aac5b7 /caffe2/image | |
| parent | e5bd6fe86d6eb07ae3544696aec3ba2e3ced8e85 (diff) | |
| download | pytorch-f61612206cee9e00011ba9945be771ddc1673092.tar.gz pytorch-f61612206cee9e00011ba9945be771ddc1673092.tar.bz2 pytorch-f61612206cee9e00011ba9945be771ddc1673092.zip | |
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
Diffstat (limited to 'caffe2/image')
| -rw-r--r-- | caffe2/image/image_input_op.h | 17 |
1 files changed, 11 insertions, 6 deletions
diff --git a/caffe2/image/image_input_op.h b/caffe2/image/image_input_op.h index 321d017847..f5db8b36dc 100644 --- a/caffe2/image/image_input_op.h +++ b/caffe2/image/image_input_op.h @@ -1243,12 +1243,14 @@ bool ImageInputOp<Context>::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<Context, CPUContext>::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<Context>::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; |