From 1ae2c1950c4ac63cff4c8c9251f3a83a9968376b Mon Sep 17 00:00:00 2001 From: Ahmed Aly Date: Thu, 28 Mar 2019 15:58:24 -0700 Subject: [PATCH] Use SetOutputTensor instead of copying outputs manually (#17770) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/17770 As title Reviewed By: dzhulgakov Differential Revision: D14370937 fbshipit-source-id: f415490c38556cf03bb13dce3643775331483448 --- caffe2/operators/inference_lstm_op.cc | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/caffe2/operators/inference_lstm_op.cc b/caffe2/operators/inference_lstm_op.cc index cfdf21f..f42fe43 100644 --- a/caffe2/operators/inference_lstm_op.cc +++ b/caffe2/operators/inference_lstm_op.cc @@ -24,21 +24,13 @@ bool InferenceLSTMOp::RunOnDevice() { bidirectional_, &context_); - std::vector allOutputs(OutputSize()); - allOutputs.at(0) = copy_ctor(std::get<0>(results)); + auto output = copy_ctor(std::get<0>(results)); if (batch_first_) { - allOutputs.at(0) = transpose(allOutputs.at(0), 0, 1, &context_); - } - allOutputs.at(1) = copy_ctor(std::get<1>(results)); - allOutputs.at(2) = copy_ctor(std::get<2>(results)); - for (int i = 0; i < OutputSize(); i++) { - auto output = XOutput(i, allOutputs.at(i).sizes(), dtype()); - context_.CopyItemsSameDevice( - allOutputs.at(i).dtype(), - allOutputs.at(i).numel(), - allOutputs.at(i).template data(), - output.template mutable_data()); + output = transpose(output, 0, 1, &context_); } + SetOutputTensor(0, copy_ctor(output)); + SetOutputTensor(1, copy_ctor(std::get<1>(results))); + SetOutputTensor(2, copy_ctor(std::get<2>(results))); return true; } -- 2.7.4