#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_);
}

REGISTER_CUDA_OPERATOR(Gather, GatherOp<CUDAContext>);
} // namespace caffe2