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

diff --git a/caffe2/image/image_input_op.h b/caffe2/image/image_input_op.h
index 321d0178473f9033231a124fa696a9e7535b32b1..f5db8b36dcef83202f6448657af13723e7b9f9fd 100644 (file)
--- 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;