Skip to content

GitLab

Unverified Commit 4e773820 authored by Wenhao Wu's avatar Wenhao Wu Committed by GitHub
Browse files

support fp16 training with gather_points op (#1758)

parent fd3a1a16
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
mmcv/ops/gather_points.py
View file @ 4e773820
...@@ -26,7 +26,7 @@ class GatherPoints(Function): ...@@ -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): ...@@ -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,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment