diff options
| author | Xiaomeng Yang <yangxm@fb.com> | 2019-02-27 12:18:52 -0800 |
|---|---|---|
| committer | Facebook Github Bot <facebook-github-bot@users.noreply.github.com> | 2019-02-27 12:34:58 -0800 |
| commit | 9709d5e787616a4dde1f0d120db951d103d1d857 (patch) | |
| tree | e77be5516fce21ba9fe52a0b6c53ae188843eee5 /caffe2/contrib | |
| parent | b4572668b48b746802580263dafa12762d202379 (diff) | |
| download | pytorch-9709d5e787616a4dde1f0d120db951d103d1d857.tar.gz pytorch-9709d5e787616a4dde1f0d120db951d103d1d857.tar.bz2 pytorch-9709d5e787616a4dde1f0d120db951d103d1d857.zip | |
Fix math::Set for large tensor (#17539)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/17539
Fix math::Set for large tensor
i-am-not-moving-c2-to-c10
Reviewed By: dzhulgakov, houseroad
Differential Revision: D14240756
fbshipit-source-id: 0ade26790be41fb26d2cc193bfa3082c7bd4e69d
Diffstat (limited to 'caffe2/contrib')
| -rw-r--r-- | caffe2/contrib/aten/aten_op.cc | 12 |
1 files changed, 7 insertions, 5 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 |