fix Expected isFloatingType error for pytorch version 1.2+

facebookresearch · Feb 7, 2020 · f610e09 · f610e09
1 parent da10b32
commit f610e09
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/pointnet2/pointnet2_utils.py b/pointnet2/pointnet2_utils.py
@@ -68,7 +68,9 @@ def forward(ctx, xyz, npoint):
         torch.Tensor
             (B, npoint) tensor containing the set
         """
-        return _ext.furthest_point_sampling(xyz, npoint)
+        fps_inds = _ext.furthest_point_sampling(xyz, npoint)
+        ctx.mark_non_differentiable(fps_inds)
+        return fps_inds
 
     @staticmethod
     def backward(xyz, a=None):
@@ -277,7 +279,9 @@ def forward(ctx, radius, nsample, xyz, new_xyz):
         torch.Tensor
             (B, npoint, nsample) tensor with the indicies of the features that form the query balls
         """
-        return _ext.ball_query(new_xyz, xyz, radius, nsample)
+        inds = _ext.ball_query(new_xyz, xyz, radius, nsample)
+        ctx.mark_non_differentiable(inds)
+        return inds
 
     @staticmethod
     def backward(ctx, a=None):