[caffe2][easy] Save heap allocation in ConcatOp (#63529)
authorScott Wolchok <swolchok@fb.com>
Wed, 1 Sep 2021 01:22:23 +0000 (18:22 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Wed, 1 Sep 2021 01:33:32 +0000 (18:33 -0700)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/63529

Output() takes an IntArrayRef, so we can just use a std::initializer_list (stack-allocated array) instead of std::vector here.
ghstack-source-id: 137085908

Test Plan: existing CI

Reviewed By: mruberry

Differential Revision: D29687400

fbshipit-source-id: 9f2a7c6679f2552c098bb1bf7befaca18e0e5d4d

caffe2/operators/concat_split_op.h

index bbe355e..f1e8f10 100644 (file)
@@ -282,7 +282,7 @@ bool ConcatOp<Context>::RunOnDevice() {
   // We can override default options(Context::GetDeviceType())
   // by explicitly passing in device type we want
   Tensor* split = Output(
-      1, std::vector<int64_t>(1, InputSize()), at::dtype<int>().device(CPU));
+      1, at::IntArrayRef({InputSize()}), at::dtype<int>().device(CPU));
   int* axis_data = split->template mutable_data<int>();
   auto& input_zero = Input(0);
   int adj_size = input_zero.dim() + (add_axis_ ? 1 : 0);