diff --git a/pytorch3d/ops/cubify.py b/pytorch3d/ops/cubify.py index 7af449b2..b47d44b0 100644 --- a/pytorch3d/ops/cubify.py +++ b/pytorch3d/ops/cubify.py @@ -223,7 +223,10 @@ def cubify(voxels, thresh, device=None, align: str = "topleft") -> Meshes: grid_faces += nyxz[:, 0].view(-1, 1) * num_verts idleverts = torch.ones(num_verts * N, dtype=torch.uint8, device=device) - idleverts.scatter_(0, grid_faces.flatten(), 0) + indices = grid_faces.flatten() + if device.type == "cpu": + indices = torch.unique(indices) + idleverts.scatter_(0, indices, 0) grid_faces -= nyxz[:, 0].view(-1, 1) * num_verts split_size = torch.bincount(nyxz[:, 0], minlength=N) faces_list = list(torch.split(grid_faces, split_size.tolist(), 0))