diff options
| -rw-r--r-- | caffe2/contrib/aten/aten_op.cc | 12 | ||||
| -rw-r--r-- | caffe2/utils/math/elementwise.cc | 24 | ||||
| -rw-r--r-- | caffe2/utils/math/elementwise.cu | 2 | ||||
| -rw-r--r-- | caffe2/utils/math/elementwise.h | 2 |
4 files changed, 21 insertions, 19 deletions
diff --git a/caffe2/contrib/aten/aten_op.cc b/caffe2/contrib/aten/aten_op.cc index 2abc25a143..aa8b7fe93e 100644 --- a/caffe2/contrib/aten/aten_op.cc +++ b/caffe2/contrib/aten/aten_op.cc @@ -4,7 +4,7 @@ namespace caffe2 { REGISTER_CPU_OPERATOR(ATen, ATenOp<CPUContext>); -template<> +template <> at::Backend ATenOp<CPUContext>::backend() const { return at::Backend::CPU; } @@ -12,14 +12,16 @@ at::Backend ATenOp<CPUContext>::backend() const { OPERATOR_SCHEMA(ATen); namespace math { + template <> void Set<at::Half, CPUContext>( - const int /*N*/, + const std::int64_t /* N */, const at::Half h, at::Half* v, CPUContext* c) { - Set(0, h.x, (uint16_t*) v, c); -} + Set(0, h.x, (uint16_t*)v, c); } -} +} // namespace math + +} // namespace caffe2 diff --git a/caffe2/utils/math/elementwise.cc b/caffe2/utils/math/elementwise.cc index 08d723c66c..96e4644b82 100644 --- a/caffe2/utils/math/elementwise.cc +++ b/caffe2/utils/math/elementwise.cc @@ -356,18 +356,18 @@ DELEGATE_SCALE(float, double, cblas_dscal) // Eigen or via custom code. //////////////////////////////////////////////////////////////////////////////// -#define CAFFE2_SPECIALIZED_SET(T) \ - template <> \ - C10_EXPORT void Set<T, CPUContext>( \ - const int N, const T alpha, T* Y, CPUContext* /* context */) { \ - if (N == 0) { \ - return; \ - } \ - if (alpha == T(0)) { \ - std::memset(Y, 0, N * sizeof(T)); \ - } else { \ - EigenVectorArrayMap<T>(Y, N).setConstant(alpha); \ - } \ +#define CAFFE2_SPECIALIZED_SET(T) \ + template <> \ + C10_EXPORT void Set<T, CPUContext>( \ + const std::int64_t N, const T alpha, T* Y, CPUContext* /* context */) { \ + if (N == 0) { \ + return; \ + } \ + if (alpha == T(0)) { \ + std::memset(Y, 0, N * sizeof(T)); \ + } else { \ + EigenVectorArrayMap<T>(Y, N).setConstant(alpha); \ + } \ } CAFFE2_SPECIALIZED_SET(float) CAFFE2_SPECIALIZED_SET(double) diff --git a/caffe2/utils/math/elementwise.cu b/caffe2/utils/math/elementwise.cu index 006fbd0b27..7509fb2429 100644 --- a/caffe2/utils/math/elementwise.cu +++ b/caffe2/utils/math/elementwise.cu @@ -34,7 +34,7 @@ __global__ void SinCosCUDAKernel(const int N, const T* X, T* S, T* C) { #define CAFFE2_SPECIALIZED_CUDA_SET(T) \ template <> \ CAFFE2_CUDA_EXPORT void Set<T, CUDAContext>( \ - const int N, const T alpha, T* Y, CUDAContext* context) { \ + const std::int64_t N, const T alpha, T* Y, CUDAContext* context) { \ if (N == 0) { \ return; \ } \ diff --git a/caffe2/utils/math/elementwise.h b/caffe2/utils/math/elementwise.h index 32ef6bee52..904853c6da 100644 --- a/caffe2/utils/math/elementwise.h +++ b/caffe2/utils/math/elementwise.h @@ -57,7 +57,7 @@ template <typename T, class Context> CAFFE2_API void Erf(int N, const T* X, T* Y, Context* context); template <typename T, class Context> -CAFFE2_API void Set(int N, T alpha, T* X, Context* context); +CAFFE2_API void Set(std::int64_t N, T alpha, T* X, Context* context); template <typename TAlpha, typename TData, class Context> CAFFE2_API void |