| 4 | #include "../utils.hpp" |
| 5 | |
| 6 | at::Tensor grouping_forward(at::Tensor features, at::Tensor indices) { |
| 7 | CHECK_CUDA(features); |
| 8 | CHECK_CUDA(indices); |
| 9 | CHECK_CONTIGUOUS(features); |
| 10 | CHECK_CONTIGUOUS(indices); |
| 11 | CHECK_IS_FLOAT(features); |
| 12 | CHECK_IS_INT(indices); |
| 13 | |
| 14 | int b = features.size(0); |
| 15 | int c = features.size(1); |
| 16 | int n = features.size(2); |
| 17 | int m = indices.size(1); |
| 18 | int u = indices.size(2); |
| 19 | at::Tensor output = torch::zeros( |
| 20 | {b, c, m, u}, at::device(features.device()).dtype(at::ScalarType::Float)); |
| 21 | grouping(b, c, n, m, u, features.data_ptr<float>(), indices.data_ptr<int>(), |
| 22 | output.data_ptr<float>()); |
| 23 | return output; |
| 24 | } |
| 25 | |
| 26 | at::Tensor grouping_backward(at::Tensor grad_y, at::Tensor indices, |
| 27 | const int n) { |
nothing calls this directly
no outgoing calls
no test coverage detected