1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
|
#include "caffe2/core/context_gpu.h"
#include "caffe2/operators/gather_op.h"
#include "caffe2/operators/gather_op.cuh"
namespace caffe2 {
template <>
bool GatherOp<CUDAContext>::RunOnDevice() {
return DispatchHelper<TensorTypes<int32_t, int64_t>>::call(
this, OperatorBase::Input<Tensor>(INDICES, CUDA));
}
template <>
template <typename Index>
bool GatherOp<CUDAContext>::DoRunWithType() {
// Use shared implementation with BatchGather
return gather_helper::gather_impl_cuda<Index>(
this, DATA, INDICES, 0, axis_, wrap_indices_, match_outer_);
}
REGISTER_CUDA_OPERATOR(Gather, GatherOp<CUDAContext>);
} // namespace caffe2
|
