mirrors/mmcv
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmcv.git synced 2025-06-03 21:54:52 +08:00
Code Issues Projects Releases Wiki Activity

support fp16 training with gather_points op (#1758)

Browse Source
This commit is contained in:
Wenhao Wu 2022-03-02 17:55:13 +08:00 committed by GitHub
parent fd3a1a16ea
commit 4e77382034
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
mmcv/ops/gather_points.py
View File

@ -26,7 +26,7 @@ class GatherPoints(Function):
B, npoint = indices.size() B, npoint = indices.size()
_, C, N = features.size() _, C, N = features.size()
output = torch.cuda.FloatTensor(B, C, npoint) output = features.new_zeros((B, C, npoint))
ext_module.gather_points_forward( ext_module.gather_points_forward(
features, indices, output, b=B, c=C, n=N, npoints=npoint) features, indices, output, b=B, c=C, n=N, npoints=npoint)
@ -41,7 +41,7 @@ class GatherPoints(Function):
idx, C, N = ctx.for_backwards idx, C, N = ctx.for_backwards
B, npoint = idx.size() B, npoint = idx.size()
grad_features = torch.cuda.FloatTensor(B, C, N).zero_() grad_features = grad_out.new_zeros((B, C, N))
grad_out_data = grad_out.data.contiguous() grad_out_data = grad_out.data.contiguous()
ext_module.gather_points_backward( ext_module.gather_points_backward(
grad_out_data, grad_out_data,