OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

[pyper] export caffe2 bucketize GPU operator to pytorch

Browse Source 
Summary: Exporting the Bucketize operator on CUDA. Also adding unit test.

Test Plan: buck test mode/dev-nosan caffe2/torch/fb/sparsenn:gpu_test -- test_bucketize

Differential Revision: D23581321

fbshipit-source-id: 7f21862984c04d840410b8718db93006f526938a
This commit is contained in:
Xiaodong Wang 2020-09-09 16:07:03 -07:00 committed by Facebook GitHub Bot
parent e0c65abd38
commit ba6ddaf04c
2 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
caffe2/operators/bucketize_op.cu
View File

@ -52,3 +52,9 @@ bool BucketizeOp<CUDAContext>::RunOnDevice() {
REGISTER_CUDA_OPERATOR(Bucketize, BucketizeOp<CUDAContext>);
} // namespace caffe2
using BucketizeCUDA = caffe2::BucketizeOp<caffe2::CUDAContext>;
C10_EXPORT_CAFFE2_OP_TO_C10_CUDA(
Bucketize,
BucketizeCUDA);

2
caffe2/operators/bucketize_op.h
View File

@ -9,7 +9,7 @@
#include "caffe2/core/operator.h"
#include "caffe2/utils/math.h"
C10_DECLARE_EXPORT_CAFFE2_OP_TO_C10(BucketizeOp);
C10_DECLARE_EXPORT_CAFFE2_OP_TO_C10(Bucketize);
namespace caffe2 {