6
6
#include " caffe2/core/types.h"
7
7
#include " caffe2/utils/math.h"
8
8
#include " caffe2/utils/string_utils.h"
9
+ #include < c10/util/accumulate.h>
9
10
10
11
namespace caffe2 {
11
12
@@ -277,13 +278,13 @@ bool SplitByLengthsOp<Context>::RunOnDevice() {
277
278
278
279
template <class Context >
279
280
bool ConcatOp<Context>::RunOnDevice() {
280
- auto * output = Output (0 );
281
+ auto * const output = Output (0 );
281
282
282
283
// We can override default options(Context::GetDeviceType())
283
284
// by explicitly passing in device type we want
284
- Tensor* split = Output (
285
+ Tensor * const split = Output (
285
286
1 , at::IntArrayRef ({InputSize ()}), at::dtype<int >().device (CPU));
286
- int * axis_data = split->template mutable_data <int >();
287
+ int * const axis_data = split->template mutable_data <int >();
287
288
auto & input_zero = Input (0 );
288
289
int adj_size = input_zero.dim () + (add_axis_ ? 1 : 0 );
289
290
int canonical_axis = canonical_axis_index_ (axis_, adj_size);
@@ -347,7 +348,13 @@ bool ConcatOp<Context>::RunOnDevice() {
347
348
} else {
348
349
output_dims[canonical_axis] = output_channels;
349
350
}
351
+
350
352
output->Resize (output_dims);
353
+ auto *const output_ptr = static_cast <char *>(output->raw_mutable_data (input_zero.dtype ()));
354
+ if (output_ptr == nullptr ){
355
+ return true ;
356
+ }
357
+
351
358
size_t output_offset = 0 ;
352
359
for (int i = 0 ; i < InputSize (); ++i) {
353
360
auto & input = Input (i);
@@ -358,8 +365,7 @@ bool ConcatOp<Context>::RunOnDevice() {
358
365
axis_dim * after,
359
366
input.raw_data (),
360
367
axis_dim * after,
361
- static_cast <char *>(output->raw_mutable_data (input_zero.dtype ())) +
362
- output_offset,
368
+ output_ptr + output_offset,
363
369
output_channels * after,
364
370
&context_,
365
371
input_zero.dtype ().copy ());
0 commit comments